Constrain User-Defined Operator

[wilberth-narvaez]

Hello,

I have a molecular wavefunction that is a linear combination of Gaussian type orbitals (GTOs).

$$ \begin{align} y(r_{1},r_{2},r_{3}) = \sum_{i=1}^{3}\sum_{j=1}^{N} c_{i,j} \times x^{a}y^{b}z^{c} \times e^{-\alpha_{i,j} \times r_{i}^{2}} \end{align} $$

I want to use PySR to find a linear combination of Slater type orbitals (STOs) that can approximate the wavefunction.

$$ \begin{align} f(r_{1},r_{2},r_{3}) = \sum_{i=1}^{3}\sum_{j=1}^{N} c_{i,j} \times r_{i,j}^{a} \times e^{-\alpha_{i,j} \times r_{i}} \end{align} $$

Because it's of practical importance to me for each STO to be centered on an individual atom, I want to know if it is possible to implement a constraint that prevents an STO from having more than a single variable times a constant as an input. In short, I want to provide the function $s(r) = re^{-r}$ as my only unary operator to find the solution of the form $f(r_{1},r_{2}) = c_{0}r_{1}e^{-\alpha_{0}r_1} + c_{1}r_{2}e^{-\alpha_{1}r_{2}}$, where $r_{1}$ and $r_{2}$ are the data's features. However, the program predominantly explores solutions of the form $f(r_{1},r_{2}) = (r_{1} \times r_{2})e^{-\alpha_{1}r_{1} \times + r_{2} + c_{0} + c_{1}}$.

Regards,
W. N.

[MilesCranmer]

Hi @narvaez-programming,
The simplest solution is to set up nested_constraints and constraints: https://astroautomata.com/PySR/api/#working-with-complexities. For example, constraints={"sto": 3} would mean that the operator sto(r) = r*exp(-r) can have at most a complexity of 3 in its argument, meaning something like sto(a * r) or simpler would be allowed.
Cheers,
Miles

[MilesCranmer]

Since you want it to be a linear combination, you could also add the constraint that "*" have at most complexity 1 in its left argument (such as a constant). For example constraints={"sto": 3, "*": (1, -1)}. Here the -1 just means any complexity is allowed in the right argument.