Coevolution with simultaneous function evaluation

[bengcohen]

Hello all and thank you for a great SR package! I am trying to use PySR to coevolve two expressions that depend on each other when one of the outputs has too few measurements (or none at all). Essentially I would like to find $f_0$ and $f_1$ such that $y_0=f_0(X)$ and $y_1=f_1(X, y_0)$. Since there is very little data for $y_0$, I would like to evaluate the fitness of both expressions together. Is there a simple way to implement this in PySR? I took a look at MulitSRRegressor, but that did not seem to work quite like this.
Thank you for your help!

[MilesCranmer]

Hi @bengcohen,

Thanks and cool idea. I think you could do this with a custom objective that creates that specific structure, like this example: https://astroautomata.com/PySR/examples/#9-custom-objectives.

e.g., say that you set the last column of X to be all 0. You can put y0 into that column during the evaluation.

tree.degree != 2 && return L(1e10)

# Split the tree at the root node
f0 = tree.l
f1 = tree.r

y0, flag0 = eval_tree_array(f0, dataset.X, options)
!flag0 && return L(1e9)  # Return a bit lower loss for this breakpoint, because at least it got this far

modified_X = copy(dataset.X)
modified_X[end, :] .= y0  # Set last column to y0

y1, flag1 = eval_tree_array(f1, modified_X, options)
!flag1 && return L(1e8)

# Now you have y0 and y1 and can do your evaluation on them

However the one tricky thing is the printed expression (and the one returned to Python) is not really the same thing as the one you are evaluating inside the loss function. So you will need to manually interpret the results afterward. But see the custom objectives example for more discussion about this.

Also if you need to fit y0 to a different ground truth than y1, what you could do is pass the y1 ground truth to the "weights" and access it with dataset.weights.

This is a bit easier on the Julia side as there is the dataset.extras parameter for storing additional variables. But it would be great to have that in the Python side as well, I just need to find time to add it.

[bengcohen]

Thanks for the thoughtful response @MilesCranmer!

Splitting the tree in this way looks like it could work well. I'll give it a try.

If you want any support adding the dataset.extras feature on the Python side, please let me know.

Thank you again!

[MilesCranmer]

Adding dataset.extras would be extremely appreciated!

Changes needed:

1. Here: https://github.com/MilesCranmer/SymbolicRegression.jl/blob/feab045092498b018b03fbf83e2cfc32543d9775/src/SymbolicRegression.jl#L357 add an argument extra_data::NamedTuple.
2. Pass extra_data to construct_datasets here: https://github.com/MilesCranmer/SymbolicRegression.jl/blob/feab045092498b018b03fbf83e2cfc32543d9775/src/SymbolicRegression.jl#L375
3. Add an extra argument to this function: https://github.com/MilesCranmer/SymbolicRegression.jl/blob/feab045092498b018b03fbf83e2cfc32543d9775/src/SearchUtils.jl#L346 (i.e., extra_data would become extra here - just want to make it extra_data at the API level so its easier to know what "extra" refers to).
4. In that function, pass extra=extra to the bottom of the Dataset call.

That should be it on the Julia side. On the Python side:

1. Add an extra_data parameter to the fit call:

   PySR/pysr/sr.py

   Line 1769 in b56466b

   y_units: Optional[List[str]] = None,

   . I think this would be a Optional[Dict[str, Any]] type.
2. I think we would basically skip all of the preprocessing since it's completely flexible what the user could pass.
3. So it should just get passed directly to self._run here:

   PySR/pysr/sr.py

   Line 1901 in b56466b

   self._run(X, y, mutated_params, weights=weights, seed=seed)

4. Add it to def _run(self, ...) as extra_data.
5. In this function, you would want to convert the dict to a named tuple. You can do this as follows:

Main._extra_data = extra_data
extra_data_named_tuple = jl.eval("(; (Symbol(k) => v for (k,v) in x)...)")

Just make sure to only do this if extra_data is not None.

7. Pass this extra_data_named_tuple to this call here:

   PySR/pysr/sr.py

   Line 1730 in b56466b

   self.raw_julia_state_ = SymbolicRegression.equation_search(

   as another line.

[bengcohen]

Great! I can certainly help implement these changes. Is there a timeline that you would like adhere to for these changes? If not, I can update you on any progress in Mid-February.

[MilesCranmer]

Sure that is good. 👍