PySR
Symbolic Regression in Python and Julia
Docs × Examples × API × Forum × Agents
-
-
-PySR searches for symbolic expressions which optimize a particular objective.
-
-https://github.com/astroautomata/PySR/assets/7593028/c8511a49-b408-488f-8f18-b1749078268f
-
-
-# PySR: High-Performance Symbolic Regression in Python and Julia
-
-| **Docs** | **Forums** | **Paper** | **colab demo** |
-|:---:|:---:|:---:|:---:|
-|[](https://ai.damtp.cam.ac.uk/pysr/)|[](https://github.com/astroautomata/PySR/discussions)|[](https://arxiv.org/abs/2305.01582)|[](https://colab.research.google.com/github/astroautomata/PySR/blob/master/examples/pysr_demo.ipynb)|
-
-| **pip** | **conda** | **Stats** |
-| :---: | :---: | :---: |
-|[](https://badge.fury.io/py/pysr)|[](https://anaconda.org/conda-forge/pysr)|
+```bash
+pip install pysr
+```
If you find PySR useful, please cite the paper [arXiv:2305.01582](https://arxiv.org/abs/2305.01582).
-If you've finished a project with PySR, please submit a PR to showcase your work on the [research showcase page](https://ai.damtp.cam.ac.uk/pysr/papers)!
+If you've finished a project with PySR, please submit a PR to showcase your work on the [research showcase page](https://pysr.ai/papers)!
**Contents**:
- [Why PySR?](#why-pysr)
- [Installation](#installation)
- [Quickstart](#quickstart)
-- [→ Documentation](https://ai.damtp.cam.ac.uk/pysr)
+- [→ Documentation](https://pysr.ai)
- [Contributors](#contributors-)
-pip: [](https://pypi.org/project/pysr/)
conda: [](https://anaconda.org/conda-forge/pysr)
|
-
-conda: [](https://anaconda.org/conda-forge/pysr)
-
-### Test status
-
-| **Linux** | **Windows** | **macOS** |
-|---|---|---|
-|[](https://github.com/astroautomata/PySR/actions/workflows/CI.yml)|[](https://github.com/astroautomata/PySR/actions/workflows/CI_Windows.yml)|[](https://github.com/astroautomata/PySR/actions/workflows/CI_mac.yml)|
-| **Docker** | **Conda** | **Coverage** |
-|[](https://github.com/astroautomata/PySR/actions/workflows/CI_docker.yml)|[](https://github.com/astroautomata/PySR/actions/workflows/CI_conda_forge.yml)|[](https://codecov.io/gh/astroautomata/PySR)|
-
-
-
## Why PySR?
PySR is an open-source tool for *Symbolic Regression*: a machine learning
@@ -65,24 +38,29 @@ an explicit and powerful way to interpret deep neural networks.
## Installation
-### Pip
-
-You can install PySR with pip:
+You can install PySR with pip or conda:
```bash
pip install pysr
```
+```bash
+conda install -c conda-forge pysr
+```
+
Julia dependencies will be installed at first import.
-### Conda
+### AI Agent Skill
-Similarly, with conda:
+Install the [PySR skill](skills/pysr/SKILL.md) to provide your coding
+agents guidance using PySR as part of a larger agentic loop:
```bash
-conda install -c conda-forge pysr
+npx skills add astroautomata/PySR
```
+Add `--global` to install the skill globally.
+
@@ -255,16 +233,16 @@ model = PySRRegressor.from_file("hall_of_fame.2022-08-10_100832.281.pkl")
There are several other useful features such as denoising (e.g., `denoise=True`),
feature selection (e.g., `select_k_features=3`).
-For examples of these and other features, see the [examples page](https://ai.damtp.cam.ac.uk/pysr/examples).
-For a detailed look at more options, see the [options page](https://ai.damtp.cam.ac.uk/pysr/options).
-You can also see the full API at [this page](https://ai.damtp.cam.ac.uk/pysr/api).
-There are also tips for tuning PySR on [this page](https://ai.damtp.cam.ac.uk/pysr/tuning).
+For examples of these and other features, see the [examples page](https://pysr.ai/examples).
+For a detailed look at more options, see the [options page](https://pysr.ai/options).
+You can also see the full API at [this page](https://pysr.ai/api).
+There are also tips for tuning PySR on [this page](https://pysr.ai/tuning).
### Detailed Example
The following code makes use of as many PySR features as possible.
Note that is just a demonstration of features and you should not use this example as-is.
-For details on what each parameter does, check out the [API page](https://ai.damtp.cam.ac.uk/pysr/api/).
+For details on what each parameter does, check out the [API page](https://pysr.ai/api).
```python
model = PySRRegressor(
@@ -379,51 +357,103 @@ If you have an idea for a new feature, don't hesitate to share it on the [issues
\">![](\"<%=)
?s=<%= options.imageSize %>\" width=\"<%= options.imageSize %>px;\" alt=\"<%= contributor.name %>\"/>
<%= contributor.name %>"
}
diff --git a/docs/all_contributors/.gitignore b/docs/all_contributors/.gitignore
index 3c3629e64..22fbdbe40 100644
--- a/docs/all_contributors/.gitignore
+++ b/docs/all_contributors/.gitignore
@@ -1 +1,3 @@
node_modules
+yarn.lock
+package-lock.json
diff --git a/docs/all_contributors/run_all_contrib.sh b/docs/all_contributors/run_all_contrib.sh
index 562efda08..278e30ed0 100755
--- a/docs/all_contributors/run_all_contrib.sh
+++ b/docs/all_contributors/run_all_contrib.sh
@@ -1,4 +1,4 @@
#!/bin/bash
-yarn install --frozen-lockfile
+yarn install
yarn all-contributors $@
diff --git a/docs/all_contributors/yarn.lock b/docs/all_contributors/yarn.lock
deleted file mode 100644
index 5a7097e39..000000000
--- a/docs/all_contributors/yarn.lock
+++ /dev/null
@@ -1,427 +0,0 @@
-# THIS IS AN AUTOGENERATED FILE. DO NOT EDIT THIS FILE DIRECTLY.
-# yarn lockfile v1
-
-
-"@babel/runtime@^7.18.9", "@babel/runtime@^7.7.6":
- version "7.22.6"
- resolved "https://registry.yarnpkg.com/@babel/runtime/-/runtime-7.22.6.tgz#57d64b9ae3cff1d67eb067ae117dac087f5bd438"
- integrity sha512-wDb5pWm4WDdF6LFUde3Jl8WzPA+3ZbxYqkC6xAXuD3irdEHN1k0NfTRrJD8ZD378SJ61miMLCqIOXYhd8x+AJQ==
- dependencies:
- regenerator-runtime "^0.13.11"
-
-all-contributors-cli@6.25.1:
- version "6.25.1"
- resolved "https://registry.yarnpkg.com/all-contributors-cli/-/all-contributors-cli-6.25.1.tgz#d6fbd4e7264939dc8efe897a9ccddc5329e46235"
- integrity sha512-Q9MfsO6FH09h71IG6OobfESA7lbooB3/bnO2wnuXbLPye1lFsYsa/jpNZP0YBx6zbhwXqqm6CXRZ9HO/tmD4NQ==
- dependencies:
- "@babel/runtime" "^7.7.6"
- async "^3.1.0"
- chalk "^4.0.0"
- didyoumean "^1.2.1"
- inquirer "^7.3.3"
- json-fixer "^1.6.8"
- lodash "^4.11.2"
- node-fetch "^2.6.0"
- pify "^5.0.0"
- yargs "^15.0.1"
-
-ansi-escapes@^4.2.1:
- version "4.3.2"
- resolved "https://registry.yarnpkg.com/ansi-escapes/-/ansi-escapes-4.3.2.tgz#6b2291d1db7d98b6521d5f1efa42d0f3a9feb65e"
- integrity sha512-gKXj5ALrKWQLsYG9jlTRmR/xKluxHV+Z9QEwNIgCfM1/uwPMCuzVVnh5mwTd+OuBZcwSIMbqssNWRm1lE51QaQ==
- dependencies:
- type-fest "^0.21.3"
-
-ansi-regex@^5.0.1:
- version "5.0.1"
- resolved "https://registry.yarnpkg.com/ansi-regex/-/ansi-regex-5.0.1.tgz#082cb2c89c9fe8659a311a53bd6a4dc5301db304"
- integrity sha512-quJQXlTSUGL2LH9SUXo8VwsY4soanhgo6LNSm84E1LBcE8s3O0wpdiRzyR9z/ZZJMlMWv37qOOb9pdJlMUEKFQ==
-
-ansi-styles@^4.0.0, ansi-styles@^4.1.0:
- version "4.3.0"
- resolved "https://registry.yarnpkg.com/ansi-styles/-/ansi-styles-4.3.0.tgz#edd803628ae71c04c85ae7a0906edad34b648937"
- integrity sha512-zbB9rCJAT1rbjiVDb2hqKFHNYLxgtk8NURxZ3IZwD3F6NtxbXZQCnnSi1Lkx+IDohdPlFp222wVALIheZJQSEg==
- dependencies:
- color-convert "^2.0.1"
-
-async@^3.1.0:
- version "3.2.4"
- resolved "https://registry.yarnpkg.com/async/-/async-3.2.4.tgz#2d22e00f8cddeb5fde5dd33522b56d1cf569a81c"
- integrity sha512-iAB+JbDEGXhyIUavoDl9WP/Jj106Kz9DEn1DPgYw5ruDn0e3Wgi3sKFm55sASdGBNOQB8F59d9qQ7deqrHA8wQ==
-
-camelcase@^5.0.0:
- version "5.3.1"
- resolved "https://registry.yarnpkg.com/camelcase/-/camelcase-5.3.1.tgz#e3c9b31569e106811df242f715725a1f4c494320"
- integrity sha512-L28STB170nwWS63UjtlEOE3dldQApaJXZkOI1uMFfzf3rRuPegHaHesyee+YxQ+W6SvRDQV6UrdOdRiR153wJg==
-
-chalk@^4.0.0, chalk@^4.1.0, chalk@^4.1.2:
- version "4.1.2"
- resolved "https://registry.yarnpkg.com/chalk/-/chalk-4.1.2.tgz#aac4e2b7734a740867aeb16bf02aad556a1e7a01"
- integrity sha512-oKnbhFyRIXpUuez8iBMmyEa4nbj4IOQyuhc/wy9kY7/WVPcwIO9VA668Pu8RkO7+0G76SLROeyw9CpQ061i4mA==
- dependencies:
- ansi-styles "^4.1.0"
- supports-color "^7.1.0"
-
-chardet@^0.7.0:
- version "0.7.0"
- resolved "https://registry.yarnpkg.com/chardet/-/chardet-0.7.0.tgz#90094849f0937f2eedc2425d0d28a9e5f0cbad9e"
- integrity sha512-mT8iDcrh03qDGRRmoA2hmBJnxpllMR+0/0qlzjqZES6NdiWDcZkCNAk4rPFZ9Q85r27unkiNNg8ZOiwZXBHwcA==
-
-cli-cursor@^3.1.0:
- version "3.1.0"
- resolved "https://registry.yarnpkg.com/cli-cursor/-/cli-cursor-3.1.0.tgz#264305a7ae490d1d03bf0c9ba7c925d1753af307"
- integrity sha512-I/zHAwsKf9FqGoXM4WWRACob9+SNukZTd94DWF57E4toouRulbCxcUh6RKUEOQlYTHJnzkPMySvPNaaSLNfLZw==
- dependencies:
- restore-cursor "^3.1.0"
-
-cli-width@^3.0.0:
- version "3.0.0"
- resolved "https://registry.yarnpkg.com/cli-width/-/cli-width-3.0.0.tgz#a2f48437a2caa9a22436e794bf071ec9e61cedf6"
- integrity sha512-FxqpkPPwu1HjuN93Omfm4h8uIanXofW0RxVEW3k5RKx+mJJYSthzNhp32Kzxxy3YAEZ/Dc/EWN1vZRY0+kOhbw==
-
-cliui@^6.0.0:
- version "6.0.0"
- resolved "https://registry.yarnpkg.com/cliui/-/cliui-6.0.0.tgz#511d702c0c4e41ca156d7d0e96021f23e13225b1"
- integrity sha512-t6wbgtoCXvAzst7QgXxJYqPt0usEfbgQdftEPbLL/cvv6HPE5VgvqCuAIDR0NgU52ds6rFwqrgakNLrHEjCbrQ==
- dependencies:
- string-width "^4.2.0"
- strip-ansi "^6.0.0"
- wrap-ansi "^6.2.0"
-
-color-convert@^2.0.1:
- version "2.0.1"
- resolved "https://registry.yarnpkg.com/color-convert/-/color-convert-2.0.1.tgz#72d3a68d598c9bdb3af2ad1e84f21d896abd4de3"
- integrity sha512-RRECPsj7iu/xb5oKYcsFHSppFNnsj/52OVTRKb4zP5onXwVF3zVmmToNcOfGC+CRDpfK/U584fMg38ZHCaElKQ==
- dependencies:
- color-name "~1.1.4"
-
-color-name@~1.1.4:
- version "1.1.4"
- resolved "https://registry.yarnpkg.com/color-name/-/color-name-1.1.4.tgz#c2a09a87acbde69543de6f63fa3995c826c536a2"
- integrity sha512-dOy+3AuW3a2wNbZHIuMZpTcgjGuLU/uBL/ubcZF9OXbDo8ff4O8yVp5Bf0efS8uEoYo5q4Fx7dY9OgQGXgAsQA==
-
-decamelize@^1.2.0:
- version "1.2.0"
- resolved "https://registry.yarnpkg.com/decamelize/-/decamelize-1.2.0.tgz#f6534d15148269b20352e7bee26f501f9a191290"
- integrity sha512-z2S+W9X73hAUUki+N+9Za2lBlun89zigOyGrsax+KUQ6wKW4ZoWpEYBkGhQjwAjjDCkWxhY0VKEhk8wzY7F5cA==
-
-didyoumean@^1.2.1:
- version "1.2.2"
- resolved "https://registry.yarnpkg.com/didyoumean/-/didyoumean-1.2.2.tgz#989346ffe9e839b4555ecf5666edea0d3e8ad037"
- integrity sha512-gxtyfqMg7GKyhQmb056K7M3xszy/myH8w+B4RT+QXBQsvAOdc3XymqDDPHx1BgPgsdAA5SIifona89YtRATDzw==
-
-emoji-regex@^8.0.0:
- version "8.0.0"
- resolved "https://registry.yarnpkg.com/emoji-regex/-/emoji-regex-8.0.0.tgz#e818fd69ce5ccfcb404594f842963bf53164cc37"
- integrity sha512-MSjYzcWNOA0ewAHpz0MxpYFvwg6yjy1NG3xteoqz644VCo/RPgnr1/GGt+ic3iJTzQ8Eu3TdM14SawnVUmGE6A==
-
-escape-string-regexp@^1.0.5:
- version "1.0.5"
- resolved "https://registry.yarnpkg.com/escape-string-regexp/-/escape-string-regexp-1.0.5.tgz#1b61c0562190a8dff6ae3bb2cf0200ca130b86d4"
- integrity sha512-vbRorB5FUQWvla16U8R/qgaFIya2qGzwDrNmCZuYKrbdSUMG6I1ZCGQRefkRVhuOkIGVne7BQ35DSfo1qvJqFg==
-
-external-editor@^3.0.3:
- version "3.1.0"
- resolved "https://registry.yarnpkg.com/external-editor/-/external-editor-3.1.0.tgz#cb03f740befae03ea4d283caed2741a83f335495"
- integrity sha512-hMQ4CX1p1izmuLYyZqLMO/qGNw10wSv9QDCPfzXfyFrOaCSSoRfqE1Kf1s5an66J5JZC62NewG+mK49jOCtQew==
- dependencies:
- chardet "^0.7.0"
- iconv-lite "^0.4.24"
- tmp "^0.0.33"
-
-figures@^3.0.0:
- version "3.2.0"
- resolved "https://registry.yarnpkg.com/figures/-/figures-3.2.0.tgz#625c18bd293c604dc4a8ddb2febf0c88341746af"
- integrity sha512-yaduQFRKLXYOGgEn6AZau90j3ggSOyiqXU0F9JZfeXYhNa+Jk4X+s45A2zg5jns87GAFa34BBm2kXw4XpNcbdg==
- dependencies:
- escape-string-regexp "^1.0.5"
-
-find-up@^4.1.0:
- version "4.1.0"
- resolved "https://registry.yarnpkg.com/find-up/-/find-up-4.1.0.tgz#97afe7d6cdc0bc5928584b7c8d7b16e8a9aa5d19"
- integrity sha512-PpOwAdQ/YlXQ2vj8a3h8IipDuYRi3wceVQQGYWxNINccq40Anw7BlsEXCMbt1Zt+OLA6Fq9suIpIWD0OsnISlw==
- dependencies:
- locate-path "^5.0.0"
- path-exists "^4.0.0"
-
-get-caller-file@^2.0.1:
- version "2.0.5"
- resolved "https://registry.yarnpkg.com/get-caller-file/-/get-caller-file-2.0.5.tgz#4f94412a82db32f36e3b0b9741f8a97feb031f7e"
- integrity sha512-DyFP3BM/3YHTQOCUL/w0OZHR0lpKeGrxotcHWcqNEdnltqFwXVfhEBQ94eIo34AfQpo0rGki4cyIiftY06h2Fg==
-
-has-flag@^4.0.0:
- version "4.0.0"
- resolved "https://registry.yarnpkg.com/has-flag/-/has-flag-4.0.0.tgz#944771fd9c81c81265c4d6941860da06bb59479b"
- integrity sha512-EykJT/Q1KjTWctppgIAgfSO0tKVuZUjhgMr17kqTumMl6Afv3EISleU7qZUzoXDFTAHTDC4NOoG/ZxU3EvlMPQ==
-
-iconv-lite@^0.4.24:
- version "0.4.24"
- resolved "https://registry.yarnpkg.com/iconv-lite/-/iconv-lite-0.4.24.tgz#2022b4b25fbddc21d2f524974a474aafe733908b"
- integrity sha512-v3MXnZAcvnywkTUEZomIActle7RXXeedOR31wwl7VlyoXO4Qi9arvSenNQWne1TcRwhCL1HwLI21bEqdpj8/rA==
- dependencies:
- safer-buffer ">= 2.1.2 < 3"
-
-inquirer@^7.3.3:
- version "7.3.3"
- resolved "https://registry.yarnpkg.com/inquirer/-/inquirer-7.3.3.tgz#04d176b2af04afc157a83fd7c100e98ee0aad003"
- integrity sha512-JG3eIAj5V9CwcGvuOmoo6LB9kbAYT8HXffUl6memuszlwDC/qvFAJw49XJ5NROSFNPxp3iQg1GqkFhaY/CR0IA==
- dependencies:
- ansi-escapes "^4.2.1"
- chalk "^4.1.0"
- cli-cursor "^3.1.0"
- cli-width "^3.0.0"
- external-editor "^3.0.3"
- figures "^3.0.0"
- lodash "^4.17.19"
- mute-stream "0.0.8"
- run-async "^2.4.0"
- rxjs "^6.6.0"
- string-width "^4.1.0"
- strip-ansi "^6.0.0"
- through "^2.3.6"
-
-is-fullwidth-code-point@^3.0.0:
- version "3.0.0"
- resolved "https://registry.yarnpkg.com/is-fullwidth-code-point/-/is-fullwidth-code-point-3.0.0.tgz#f116f8064fe90b3f7844a38997c0b75051269f1d"
- integrity sha512-zymm5+u+sCsSWyD9qNaejV3DFvhCKclKdizYaJUuHA83RLjb7nSuGnddCHGv0hk+KY7BMAlsWeK4Ueg6EV6XQg==
-
-json-fixer@^1.6.8:
- version "1.6.15"
- resolved "https://registry.yarnpkg.com/json-fixer/-/json-fixer-1.6.15.tgz#f1f03b6771fcb383695d458c53e50b10999fba7f"
- integrity sha512-TuDuZ5KrgyjoCIppdPXBMqiGfota55+odM+j2cQ5rt/XKyKmqGB3Whz1F8SN8+60yYGy/Nu5lbRZ+rx8kBIvBw==
- dependencies:
- "@babel/runtime" "^7.18.9"
- chalk "^4.1.2"
- pegjs "^0.10.0"
-
-locate-path@^5.0.0:
- version "5.0.0"
- resolved "https://registry.yarnpkg.com/locate-path/-/locate-path-5.0.0.tgz#1afba396afd676a6d42504d0a67a3a7eb9f62aa0"
- integrity sha512-t7hw9pI+WvuwNJXwk5zVHpyhIqzg2qTlklJOf0mVxGSbe3Fp2VieZcduNYjaLDoy6p9uGpQEGWG87WpMKlNq8g==
- dependencies:
- p-locate "^4.1.0"
-
-lodash@^4.11.2, lodash@^4.17.19:
- version "4.17.21"
- resolved "https://registry.yarnpkg.com/lodash/-/lodash-4.17.21.tgz#679591c564c3bffaae8454cf0b3df370c3d6911c"
- integrity sha512-v2kDEe57lecTulaDIuNTPy3Ry4gLGJ6Z1O3vE1krgXZNrsQ+LFTGHVxVjcXPs17LhbZVGedAJv8XZ1tvj5FvSg==
-
-mimic-fn@^2.1.0:
- version "2.1.0"
- resolved "https://registry.yarnpkg.com/mimic-fn/-/mimic-fn-2.1.0.tgz#7ed2c2ccccaf84d3ffcb7a69b57711fc2083401b"
- integrity sha512-OqbOk5oEQeAZ8WXWydlu9HJjz9WVdEIvamMCcXmuqUYjTknH/sqsWvhQ3vgwKFRR1HpjvNBKQ37nbJgYzGqGcg==
-
-mute-stream@0.0.8:
- version "0.0.8"
- resolved "https://registry.yarnpkg.com/mute-stream/-/mute-stream-0.0.8.tgz#1630c42b2251ff81e2a283de96a5497ea92e5e0d"
- integrity sha512-nnbWWOkoWyUsTjKrhgD0dcz22mdkSnpYqbEjIm2nhwhuxlSkpywJmBo8h0ZqJdkp73mb90SssHkN4rsRaBAfAA==
-
-node-fetch@^2.6.0:
- version "2.6.12"
- resolved "https://registry.yarnpkg.com/node-fetch/-/node-fetch-2.6.12.tgz#02eb8e22074018e3d5a83016649d04df0e348fba"
- integrity sha512-C/fGU2E8ToujUivIO0H+tpQ6HWo4eEmchoPIoXtxCrVghxdKq+QOHqEZW7tuP3KlV3bC8FRMO5nMCC7Zm1VP6g==
- dependencies:
- whatwg-url "^5.0.0"
-
-onetime@^5.1.0:
- version "5.1.2"
- resolved "https://registry.yarnpkg.com/onetime/-/onetime-5.1.2.tgz#d0e96ebb56b07476df1dd9c4806e5237985ca45e"
- integrity sha512-kbpaSSGJTWdAY5KPVeMOKXSrPtr8C8C7wodJbcsd51jRnmD+GZu8Y0VoU6Dm5Z4vWr0Ig/1NKuWRKf7j5aaYSg==
- dependencies:
- mimic-fn "^2.1.0"
-
-os-tmpdir@~1.0.2:
- version "1.0.2"
- resolved "https://registry.yarnpkg.com/os-tmpdir/-/os-tmpdir-1.0.2.tgz#bbe67406c79aa85c5cfec766fe5734555dfa1274"
- integrity sha512-D2FR03Vir7FIu45XBY20mTb+/ZSWB00sjU9jdQXt83gDrI4Ztz5Fs7/yy74g2N5SVQY4xY1qDr4rNddwYRVX0g==
-
-p-limit@^2.2.0:
- version "2.3.0"
- resolved "https://registry.yarnpkg.com/p-limit/-/p-limit-2.3.0.tgz#3dd33c647a214fdfffd835933eb086da0dc21db1"
- integrity sha512-//88mFWSJx8lxCzwdAABTJL2MyWB12+eIY7MDL2SqLmAkeKU9qxRvWuSyTjm3FUmpBEMuFfckAIqEaVGUDxb6w==
- dependencies:
- p-try "^2.0.0"
-
-p-locate@^4.1.0:
- version "4.1.0"
- resolved "https://registry.yarnpkg.com/p-locate/-/p-locate-4.1.0.tgz#a3428bb7088b3a60292f66919278b7c297ad4f07"
- integrity sha512-R79ZZ/0wAxKGu3oYMlz8jy/kbhsNrS7SKZ7PxEHBgJ5+F2mtFW2fK2cOtBh1cHYkQsbzFV7I+EoRKe6Yt0oK7A==
- dependencies:
- p-limit "^2.2.0"
-
-p-try@^2.0.0:
- version "2.2.0"
- resolved "https://registry.yarnpkg.com/p-try/-/p-try-2.2.0.tgz#cb2868540e313d61de58fafbe35ce9004d5540e6"
- integrity sha512-R4nPAVTAU0B9D35/Gk3uJf/7XYbQcyohSKdvAxIRSNghFl4e71hVoGnBNQz9cWaXxO2I10KTC+3jMdvvoKw6dQ==
-
-path-exists@^4.0.0:
- version "4.0.0"
- resolved "https://registry.yarnpkg.com/path-exists/-/path-exists-4.0.0.tgz#513bdbe2d3b95d7762e8c1137efa195c6c61b5b3"
- integrity sha512-ak9Qy5Q7jYb2Wwcey5Fpvg2KoAc/ZIhLSLOSBmRmygPsGwkVVt0fZa0qrtMz+m6tJTAHfZQ8FnmB4MG4LWy7/w==
-
-pegjs@^0.10.0:
- version "0.10.0"
- resolved "https://registry.yarnpkg.com/pegjs/-/pegjs-0.10.0.tgz#cf8bafae6eddff4b5a7efb185269eaaf4610ddbd"
- integrity sha512-qI5+oFNEGi3L5HAxDwN2LA4Gg7irF70Zs25edhjld9QemOgp0CbvMtbFcMvFtEo1OityPrcCzkQFB8JP/hxgow==
-
-pify@^5.0.0:
- version "5.0.0"
- resolved "https://registry.yarnpkg.com/pify/-/pify-5.0.0.tgz#1f5eca3f5e87ebec28cc6d54a0e4aaf00acc127f"
- integrity sha512-eW/gHNMlxdSP6dmG6uJip6FXN0EQBwm2clYYd8Wul42Cwu/DK8HEftzsapcNdYe2MfLiIwZqsDk2RDEsTE79hA==
-
-regenerator-runtime@^0.13.11:
- version "0.13.11"
- resolved "https://registry.yarnpkg.com/regenerator-runtime/-/regenerator-runtime-0.13.11.tgz#f6dca3e7ceec20590d07ada785636a90cdca17f9"
- integrity sha512-kY1AZVr2Ra+t+piVaJ4gxaFaReZVH40AKNo7UCX6W+dEwBo/2oZJzqfuN1qLq1oL45o56cPaTXELwrTh8Fpggg==
-
-require-directory@^2.1.1:
- version "2.1.1"
- resolved "https://registry.yarnpkg.com/require-directory/-/require-directory-2.1.1.tgz#8c64ad5fd30dab1c976e2344ffe7f792a6a6df42"
- integrity sha512-fGxEI7+wsG9xrvdjsrlmL22OMTTiHRwAMroiEeMgq8gzoLC/PQr7RsRDSTLUg/bZAZtF+TVIkHc6/4RIKrui+Q==
-
-require-main-filename@^2.0.0:
- version "2.0.0"
- resolved "https://registry.yarnpkg.com/require-main-filename/-/require-main-filename-2.0.0.tgz#d0b329ecc7cc0f61649f62215be69af54aa8989b"
- integrity sha512-NKN5kMDylKuldxYLSUfrbo5Tuzh4hd+2E8NPPX02mZtn1VuREQToYe/ZdlJy+J3uCpfaiGF05e7B8W0iXbQHmg==
-
-restore-cursor@^3.1.0:
- version "3.1.0"
- resolved "https://registry.yarnpkg.com/restore-cursor/-/restore-cursor-3.1.0.tgz#39f67c54b3a7a58cea5236d95cf0034239631f7e"
- integrity sha512-l+sSefzHpj5qimhFSE5a8nufZYAM3sBSVMAPtYkmC+4EH2anSGaEMXSD0izRQbu9nfyQ9y5JrVmp7E8oZrUjvA==
- dependencies:
- onetime "^5.1.0"
- signal-exit "^3.0.2"
-
-run-async@^2.4.0:
- version "2.4.1"
- resolved "https://registry.yarnpkg.com/run-async/-/run-async-2.4.1.tgz#8440eccf99ea3e70bd409d49aab88e10c189a455"
- integrity sha512-tvVnVv01b8c1RrA6Ep7JkStj85Guv/YrMcwqYQnwjsAS2cTmmPGBBjAjpCW7RrSodNSoE2/qg9O4bceNvUuDgQ==
-
-rxjs@^6.6.0:
- version "6.6.7"
- resolved "https://registry.yarnpkg.com/rxjs/-/rxjs-6.6.7.tgz#90ac018acabf491bf65044235d5863c4dab804c9"
- integrity sha512-hTdwr+7yYNIT5n4AMYp85KA6yw2Va0FLa3Rguvbpa4W3I5xynaBZo41cM3XM+4Q6fRMj3sBYIR1VAmZMXYJvRQ==
- dependencies:
- tslib "^1.9.0"
-
-"safer-buffer@>= 2.1.2 < 3":
- version "2.1.2"
- resolved "https://registry.yarnpkg.com/safer-buffer/-/safer-buffer-2.1.2.tgz#44fa161b0187b9549dd84bb91802f9bd8385cd6a"
- integrity sha512-YZo3K82SD7Riyi0E1EQPojLz7kpepnSQI9IyPbHHg1XXXevb5dJI7tpyN2ADxGcQbHG7vcyRHk0cbwqcQriUtg==
-
-set-blocking@^2.0.0:
- version "2.0.0"
- resolved "https://registry.yarnpkg.com/set-blocking/-/set-blocking-2.0.0.tgz#045f9782d011ae9a6803ddd382b24392b3d890f7"
- integrity sha512-KiKBS8AnWGEyLzofFfmvKwpdPzqiy16LvQfK3yv/fVH7Bj13/wl3JSR1J+rfgRE9q7xUJK4qvgS8raSOeLUehw==
-
-signal-exit@^3.0.2:
- version "3.0.7"
- resolved "https://registry.yarnpkg.com/signal-exit/-/signal-exit-3.0.7.tgz#a9a1767f8af84155114eaabd73f99273c8f59ad9"
- integrity sha512-wnD2ZE+l+SPC/uoS0vXeE9L1+0wuaMqKlfz9AMUo38JsyLSBWSFcHR1Rri62LZc12vLr1gb3jl7iwQhgwpAbGQ==
-
-string-width@^4.1.0, string-width@^4.2.0:
- version "4.2.3"
- resolved "https://registry.yarnpkg.com/string-width/-/string-width-4.2.3.tgz#269c7117d27b05ad2e536830a8ec895ef9c6d010"
- integrity sha512-wKyQRQpjJ0sIp62ErSZdGsjMJWsap5oRNihHhu6G7JVO/9jIB6UyevL+tXuOqrng8j/cxKTWyWUwvSTriiZz/g==
- dependencies:
- emoji-regex "^8.0.0"
- is-fullwidth-code-point "^3.0.0"
- strip-ansi "^6.0.1"
-
-strip-ansi@^6.0.0, strip-ansi@^6.0.1:
- version "6.0.1"
- resolved "https://registry.yarnpkg.com/strip-ansi/-/strip-ansi-6.0.1.tgz#9e26c63d30f53443e9489495b2105d37b67a85d9"
- integrity sha512-Y38VPSHcqkFrCpFnQ9vuSXmquuv5oXOKpGeT6aGrr3o3Gc9AlVa6JBfUSOCnbxGGZF+/0ooI7KrPuUSztUdU5A==
- dependencies:
- ansi-regex "^5.0.1"
-
-supports-color@^7.1.0:
- version "7.2.0"
- resolved "https://registry.yarnpkg.com/supports-color/-/supports-color-7.2.0.tgz#1b7dcdcb32b8138801b3e478ba6a51caa89648da"
- integrity sha512-qpCAvRl9stuOHveKsn7HncJRvv501qIacKzQlO/+Lwxc9+0q2wLyv4Dfvt80/DPn2pqOBsJdDiogXGR9+OvwRw==
- dependencies:
- has-flag "^4.0.0"
-
-through@^2.3.6:
- version "2.3.8"
- resolved "https://registry.yarnpkg.com/through/-/through-2.3.8.tgz#0dd4c9ffaabc357960b1b724115d7e0e86a2e1f5"
- integrity sha512-w89qg7PI8wAdvX60bMDP+bFoD5Dvhm9oLheFp5O4a2QF0cSBGsBX4qZmadPMvVqlLJBBci+WqGGOAPvcDeNSVg==
-
-tmp@^0.0.33:
- version "0.0.33"
- resolved "https://registry.yarnpkg.com/tmp/-/tmp-0.0.33.tgz#6d34335889768d21b2bcda0aa277ced3b1bfadf9"
- integrity sha512-jRCJlojKnZ3addtTOjdIqoRuPEKBvNXcGYqzO6zWZX8KfKEpnGY5jfggJQ3EjKuu8D4bJRr0y+cYJFmYbImXGw==
- dependencies:
- os-tmpdir "~1.0.2"
-
-tr46@~0.0.3:
- version "0.0.3"
- resolved "https://registry.yarnpkg.com/tr46/-/tr46-0.0.3.tgz#8184fd347dac9cdc185992f3a6622e14b9d9ab6a"
- integrity sha512-N3WMsuqV66lT30CrXNbEjx4GEwlow3v6rr4mCcv6prnfwhS01rkgyFdjPNBYd9br7LpXV1+Emh01fHnq2Gdgrw==
-
-tslib@^1.9.0:
- version "1.14.1"
- resolved "https://registry.yarnpkg.com/tslib/-/tslib-1.14.1.tgz#cf2d38bdc34a134bcaf1091c41f6619e2f672d00"
- integrity sha512-Xni35NKzjgMrwevysHTCArtLDpPvye8zV/0E4EyYn43P7/7qvQwPh9BGkHewbMulVntbigmcT7rdX3BNo9wRJg==
-
-type-fest@^0.21.3:
- version "0.21.3"
- resolved "https://registry.yarnpkg.com/type-fest/-/type-fest-0.21.3.tgz#d260a24b0198436e133fa26a524a6d65fa3b2e37"
- integrity sha512-t0rzBq87m3fVcduHDUFhKmyyX+9eo6WQjZvf51Ea/M0Q7+T374Jp1aUiyUl0GKxp8M/OETVHSDvmkyPgvX+X2w==
-
-webidl-conversions@^3.0.0:
- version "3.0.1"
- resolved "https://registry.yarnpkg.com/webidl-conversions/-/webidl-conversions-3.0.1.tgz#24534275e2a7bc6be7bc86611cc16ae0a5654871"
- integrity sha512-2JAn3z8AR6rjK8Sm8orRC0h/bcl/DqL7tRPdGZ4I1CjdF+EaMLmYxBHyXuKL849eucPFhvBoxMsflfOb8kxaeQ==
-
-whatwg-url@^5.0.0:
- version "5.0.0"
- resolved "https://registry.yarnpkg.com/whatwg-url/-/whatwg-url-5.0.0.tgz#966454e8765462e37644d3626f6742ce8b70965d"
- integrity sha512-saE57nupxk6v3HY35+jzBwYa0rKSy0XR8JSxZPwgLr7ys0IBzhGviA1/TUGJLmSVqs8pb9AnvICXEuOHLprYTw==
- dependencies:
- tr46 "~0.0.3"
- webidl-conversions "^3.0.0"
-
-which-module@^2.0.0:
- version "2.0.1"
- resolved "https://registry.yarnpkg.com/which-module/-/which-module-2.0.1.tgz#776b1fe35d90aebe99e8ac15eb24093389a4a409"
- integrity sha512-iBdZ57RDvnOR9AGBhML2vFZf7h8vmBjhoaZqODJBFWHVtKkDmKuHai3cx5PgVMrX5YDNp27AofYbAwctSS+vhQ==
-
-wrap-ansi@^6.2.0:
- version "6.2.0"
- resolved "https://registry.yarnpkg.com/wrap-ansi/-/wrap-ansi-6.2.0.tgz#e9393ba07102e6c91a3b221478f0257cd2856e53"
- integrity sha512-r6lPcBGxZXlIcymEu7InxDMhdW0KDxpLgoFLcguasxCaJ/SOIZwINatK9KY/tf+ZrlywOKU0UDj3ATXUBfxJXA==
- dependencies:
- ansi-styles "^4.0.0"
- string-width "^4.1.0"
- strip-ansi "^6.0.0"
-
-y18n@^4.0.0:
- version "4.0.3"
- resolved "https://registry.yarnpkg.com/y18n/-/y18n-4.0.3.tgz#b5f259c82cd6e336921efd7bfd8bf560de9eeedf"
- integrity sha512-JKhqTOwSrqNA1NY5lSztJ1GrBiUodLMmIZuLiDaMRJ+itFd+ABVE8XBjOvIWL+rSqNDC74LCSFmlb/U4UZ4hJQ==
-
-yargs-parser@^18.1.2:
- version "18.1.3"
- resolved "https://registry.yarnpkg.com/yargs-parser/-/yargs-parser-18.1.3.tgz#be68c4975c6b2abf469236b0c870362fab09a7b0"
- integrity sha512-o50j0JeToy/4K6OZcaQmW6lyXXKhq7csREXcDwk2omFPJEwUNOVtJKvmDr9EI1fAJZUyZcRF7kxGBWmRXudrCQ==
- dependencies:
- camelcase "^5.0.0"
- decamelize "^1.2.0"
-
-yargs@^15.0.1:
- version "15.4.1"
- resolved "https://registry.yarnpkg.com/yargs/-/yargs-15.4.1.tgz#0d87a16de01aee9d8bec2bfbf74f67851730f4f8"
- integrity sha512-aePbxDmcYW++PaqBsJ+HYUFwCdv4LVvdnhBy78E57PIor8/OVvhMrADFFEDh8DHDFRv/O9i3lPhsENjO7QX0+A==
- dependencies:
- cliui "^6.0.0"
- decamelize "^1.2.0"
- find-up "^4.1.0"
- get-caller-file "^2.0.1"
- require-directory "^2.1.1"
- require-main-filename "^2.0.0"
- set-blocking "^2.0.0"
- string-width "^4.2.0"
- which-module "^2.0.0"
- y18n "^4.0.0"
- yargs-parser "^18.1.2"
diff --git a/docs/deploy.jl b/docs/deploy.jl
index bb051270c..d6738f1e3 100644
--- a/docs/deploy.jl
+++ b/docs/deploy.jl
@@ -1,28 +1,56 @@
#!/usr/bin/env julia
#
-# Deploy built VitePress documentation using DocumenterVitepress.jl
-# This script is called after VitePress build is complete and handles:
-# - Regular deployments to gh-pages
-# - PR preview deployments to gh-pages/previews/PR##/
-# - Dual deployment to secondary repository (ai.damtp.cam.ac.uk)
-#
+# Build and deploy VitePress documentation to versioned gh-pages directories.
+# PR previews deploy to gh-pages/previews/PR##/.
using DocumenterVitepress
-# Get deployment target from environment (for dual deployment)
-deployment_target = get(ENV, "DEPLOYMENT_TARGET", "primary")
-
println("Starting DocumenterVitepress deployment...")
-println("Deployment target: $deployment_target")
println("Event: $(get(ENV, "GITHUB_EVENT_NAME", "unknown"))")
println("Ref: $(get(ENV, "GITHUB_REF", "unknown"))")
# Get deployment decision from Documenter to determine correct subfolder
using Documenter
+include("deploy_versions.jl")
+
+struct StableRedirectVersion
+ base_version::DocumenterVitepress.BaseVersion
+ root_stubs::Dict{String,String}
+end
+
+function Documenter.determine_deploy_subfolder(deploy_decision, ::StableRedirectVersion)
+ return nothing
+end
+
+function Documenter.postprocess_before_push(
+ versions::StableRedirectVersion;
+ subfolder,
+ devurl,
+ deploy_dir,
+ dirname,
+)
+ Documenter.postprocess_before_push(
+ versions.base_version; subfolder, devurl, deploy_dir, dirname
+ )
+ root = stable_deploy_root(deploy_dir)
+ for (path, content) in versions.root_stubs
+ destination = joinpath(root, path)
+ mkpath(Base.dirname(destination))
+ write(destination, content)
+ end
+ for path in stale_root_stubs(root, keys(versions.root_stubs))
+ println("Removing stale redirect $path")
+ rm(joinpath(root, path))
+ directory = joinpath(root, Base.dirname(path))
+ while directory != root && isempty(readdir(directory))
+ rm(directory)
+ directory = Base.dirname(directory)
+ end
+ end
+ return
+end
-# Custom DeployConfig that bypasses PR origin check for cross-repo deployments
-# This allows deploying PR previews to ai.damtp.cam.ac.uk/pysr even though
-# PRs exist in astroautomata/PySR (Documenter's security check would normally block this)
+# Custom DeployConfig that bypasses PR origin checks for cross-repo previews.
struct BypassPRCheckConfig <: Documenter.DeployConfig end
function Documenter.deploy_folder(
@@ -102,161 +130,89 @@ end
Documenter.authentication_method(::BypassPRCheckConfig) = Documenter.SSH
Documenter.documenter_key(::BypassPRCheckConfig) = ENV["DOCUMENTER_KEY"]
-# Configure deployment based on target
-if deployment_target == "secondary"
- # Secondary: Use custom config to bypass PR origin check
- deploy_config = BypassPRCheckConfig()
- damtp_key = get(ENV, "DAMTP_DEPLOY_KEY", "")
- if isempty(damtp_key)
- error("DAMTP_DEPLOY_KEY environment variable is required for secondary deployment but is not set")
- end
- ENV["DOCUMENTER_KEY"] = damtp_key
+deploy_config = BypassPRCheckConfig()
+damtp_key = get(ENV, "DAMTP_DEPLOY_KEY", "")
+isempty(damtp_key) && error("DAMTP_DEPLOY_KEY environment variable is required for deployment but is not set")
+ENV["DOCUMENTER_KEY"] = damtp_key
+
+deploy_decision = Documenter.deploy_folder(
+ deploy_config;
+ repo="github.com/ai-damtp-cam-ac-uk/pysr",
+ devbranch="master",
+ devurl="dev",
+ push_preview=true,
+)
- deploy_decision = Documenter.deploy_folder(
- deploy_config;
- repo="github.com/ai-damtp-cam-ac-uk/pysr",
- devbranch="master",
- devurl="dev",
- push_preview=true,
- )
-else
- # Primary: Use normal Documenter flow with security checks
- deploy_config = Documenter.auto_detect_deploy_system()
+println("Deploy decision: all_ok=$(deploy_decision.all_ok), is_preview=$(deploy_decision.is_preview), subfolder=$(deploy_decision.subfolder)")
- deploy_decision = Documenter.deploy_folder(
- deploy_config;
- repo="github.com/astroautomata/PySR",
- devbranch="master",
- devurl="dev",
- push_preview=true,
- )
+if !deploy_decision.all_ok || isempty(deploy_decision.subfolder)
+ println("Deployment skipped because no deployable subfolder was selected")
+ exit(0)
end
-println("Deploy decision: all_ok=$(deploy_decision.all_ok), is_preview=$(deploy_decision.is_preview), subfolder=$(deploy_decision.subfolder)")
-
-# Build VitePress with the correct base path for this deployment
-# VitePress needs the base path set at build time (it's hardcoded into assets)
-# Primary uses /PySR/ (capital P), secondary uses /pysr/ (lowercase p)
-base_prefix = deployment_target == "secondary" ? "/pysr/" : "/PySR/"
-full_base = "$(base_prefix)$(deploy_decision.subfolder)$(isempty(deploy_decision.subfolder) ? "" : "/")"
+subfolder = deploy_decision.subfolder
-# The version picker needs __DEPLOY_ABSPATH__ to construct URLs to sibling versions
-# This should be the shared prefix (e.g., /pysr/ or /PySR/)
-deploy_abspath = base_prefix
+base_prefix = "/"
+repo_url = "github.com/ai-damtp-cam-ac-uk/pysr.git"
-println("Building VitePress with base: $full_base (deploy abspath: $deploy_abspath)")
+# VitePress bakes the base path into every asset URL at build time
+full_base = "$(base_prefix)$(subfolder)/"
+println("Building VitePress with base: $full_base (deploy abspath: $base_prefix)")
config_path = joinpath(@__DIR__, "src", ".vitepress", "config.mts")
original_config = read(config_path, String)
-# Match either /pysr/ or /PySR/ in the config
-modified_config = replace(original_config, r"base:\s*'/[Pp]y[Ss][Rr]/'" => "base: '$full_base'")
-# Also update __DEPLOY_ABSPATH__ so version picker works correctly
-modified_config = replace(
- modified_config,
- r"__DEPLOY_ABSPATH__\s*:\s*JSON\.stringify\(getBaseRepository\([^)]+\)\)" =>
- "__DEPLOY_ABSPATH__: JSON.stringify('$deploy_abspath')",
-)
-# Primary deployment should point to Cambridge as canonical
-# Secondary (Cambridge) should have empty canonical (it's already the canonical site)
-canonical_domain = deployment_target == "primary" ? "https://ai.damtp.cam.ac.uk/pysr/" : ""
-modified_config = replace(
- modified_config,
- r"const canonicalDomain = '';" =>
- "const canonicalDomain = '$canonical_domain';",
-)
+modified_config = replace(original_config, r"base:\s*'/'" => "base: '$full_base'")
write(config_path, modified_config)
try
- # Build VitePress (outputs to docs/dist/)
cd(@__DIR__) do
run(`npm run build:vitepress`)
end
println("VitePress build complete")
finally
- # Restore original config (don't commit the modified version)
write(config_path, original_config)
println("Restored original config.mts")
end
-# DocumenterVitepress expects files in dist/1/ (versioned subdirectory)
-# But VitePress builds directly to dist/, so we need to restructure
-dist_root = joinpath(@__DIR__, "dist")
-dist_versioned = joinpath(dist_root, "1")
-
-if !isdir(dist_versioned) && isdir(dist_root)
- println("Restructuring dist/ for DocumenterVitepress...")
- # Move all files from dist/ to dist/1/
- temp_dir = joinpath(@__DIR__, "dist_temp")
- mv(dist_root, temp_dir)
- mkpath(dist_root)
- mv(temp_dir, dist_versioned)
-end
-
-# Create bases.txt with the correct subfolder from deploy_decision
-# This tells DocumenterVitepress where to deploy (e.g., "dev", "previews/PR1056", "v1.2.3")
-# Don't overwrite if it already exists with content (DocumenterVitepress may generate multi-base configs)
-bases_file = joinpath(dist_root, "bases.txt")
-if !isfile(bases_file)
- println("Creating bases.txt with subfolder: $(deploy_decision.subfolder)")
- write(bases_file, "$(deploy_decision.subfolder)\n")
-else
- bases = filter(!isempty, readlines(bases_file))
- if isempty(bases)
- println("Fixing empty bases.txt with subfolder: $(deploy_decision.subfolder)")
- write(bases_file, "$(deploy_decision.subfolder)\n")
- else
- println("bases.txt already exists with $(length(bases)) bases: $bases")
- # Don't overwrite it - DocumenterVitepress may have generated multiple bases
- end
-end
+# The version picker resolves each folder's version through this file
+write(
+ joinpath(@__DIR__, "dist", "siteinfo.js"),
+ "var DOCUMENTER_CURRENT_VERSION = $(repr(subfolder));\n",
+)
-# Create redirect index.html at root to redirect to dev
-# Only do this when deploying to dev (not for version tags)
-# Once there are tagged versions, DocumenterVitepress will handle this automatically
-if deploy_decision.subfolder == "dev"
- redirect_html = """
-
-"""
- redirect_file = joinpath(dist_root, "index.html")
- println("Creating redirect index.html to ./dev/")
- write(redirect_file, redirect_html)
+deploy(folder, target; versions = DocumenterVitepress.BaseVersion(folder)) =
+ Documenter.deploydocs(;
+ root = @__DIR__,
+ repo = repo_url,
+ deploy_config = deploy_config,
+ push_preview = true,
+ devbranch = "master",
+ devurl = "dev",
+ target = target,
+ dirname = folder,
+ versions = versions,
+ )
- # Create siteinfo.js for dev version
- # This file is required by the version picker to show all available versions
- siteinfo_dir = joinpath(dist_versioned, "dev")
- if isdir(siteinfo_dir)
- siteinfo_content = """var DOCUMENTER_CURRENT_VERSION = "dev";
-"""
- siteinfo_file = joinpath(siteinfo_dir, "siteinfo.js")
- println("Creating siteinfo.js for dev version")
- write(siteinfo_file, siteinfo_content)
- else
- println("Warning: dev directory not found at $(siteinfo_dir), skipping siteinfo.js creation")
+deploy(subfolder, "dist")
+
+if claims_stable(subfolder)
+ println("Pointing stable at $subfolder")
+ dist_dir = joinpath(@__DIR__, "dist")
+ redirect_files = stable_redirect_files(built_pages(dist_dir), subfolder)
+ stable_dir = joinpath(@__DIR__, "dist_stable")
+ mkpath(stable_dir)
+ root_stubs = Dict{String,String}()
+ for (path, content) in redirect_files
+ if startswith(path, "stable/")
+ destination = joinpath(stable_dir, path[(length("stable/") + 1):end])
+ mkpath(Base.dirname(destination))
+ write(destination, content)
+ else
+ root_stubs[path] = content
+ end
end
-end
-
-if deployment_target == "secondary"
- # Secondary deployment to ai.damtp.cam.ac.uk
- println("Deploying to secondary repository (ai.damtp.cam.ac.uk)")
- DocumenterVitepress.deploydocs(;
- root=@__DIR__,
- repo="github.com/ai-damtp-cam-ac-uk/pysr.git",
- deploy_config=deploy_config, # Use custom config with bypassed PR check
- push_preview=true,
- target="dist",
- devbranch="master",
- )
-else
- # Primary deployment to astroautomata/PySR
- println("Deploying to primary repository (astroautomata/PySR)")
- DocumenterVitepress.deploydocs(;
- root=@__DIR__,
- repo="github.com/astroautomata/PySR.git",
- deploy_config=deploy_config, # Use normal GitHubActions config
- push_preview=true,
- target="dist",
- devbranch="master",
- )
+ versions = StableRedirectVersion(DocumenterVitepress.BaseVersion("stable"), root_stubs)
+ deploy("stable", "dist_stable"; versions)
end
println("Deployment complete!")
diff --git a/docs/deploy_versions.jl b/docs/deploy_versions.jl
new file mode 100644
index 000000000..53379019d
--- /dev/null
+++ b/docs/deploy_versions.jl
@@ -0,0 +1,80 @@
+using DocumenterVitepress
+
+# Whether this tag is the newest release, and so owns the `stable` alias.
+#
+# Julia parses Python-style prerelease tags such as v2.0.0a2 as v2.0.0-a2, and
+# determine_bases then reports the alpha as owning the v2.0.0 release folder, so only
+# plain release tags reach it.
+function claims_stable(subfolder; kwargs...)
+ occursin(r"^v\d+(\.\d+){0,2}$", subfolder) || return false
+ return "stable" in DocumenterVitepress.determine_bases(subfolder; keep = :patch, kwargs...)
+end
+
+const STUB_MARKER = ""
+
+stable_deploy_root(deploy_dir) = dirname(rstrip(deploy_dir, '/'))
+
+# Directories at the root of gh-pages that a `stable` deployment must never rewrite.
+const PRESERVED_ROOT_DIRS = (".git", "previews", "stable")
+
+redirect_stub(target) = """$(STUB_MARKER)
+
+"""
+
+# Page paths are URL-relative and may be nested, so every stub has to climb back to the
+# root before naming its target.
+climb(depth) = repeat("../", depth)
+depth_of(path) = count(==('/'), path)
+
+function add_stub!(files, path, target)
+ haskey(files, path) && error("Conflicting redirect stubs for $path")
+ files[path] = redirect_stub(target)
+ return files
+end
+
+function stable_redirect_files(page_paths, stable_version)
+ pages = sort(filter(page -> page != "index" && basename(page) != "404", page_paths))
+ files = Dict("stable/index.html" => redirect_stub("../$(stable_version)/"))
+ for page in pages
+ depth = depth_of(page)
+ add_stub!(
+ files,
+ "stable/$(page).html",
+ "$(climb(depth + 1))$(stable_version)/$(page).html",
+ )
+ add_stub!(files, "$(page).html", "$(climb(depth))stable/$(page).html")
+ # A page that is already a directory index has no separate directory form.
+ basename(page) == "index" && continue
+ add_stub!(files, "$(page)/index.html", "$(climb(depth + 1))stable/$(page).html")
+ end
+ return files
+end
+
+# Every `.html` file below `dist`, named the way its URL is, so nested pages get stubs.
+function built_pages(dist_dir)
+ pages = String[]
+ for (directory, _, files) in walkdir(dist_dir), file in files
+ endswith(file, ".html") || continue
+ push!(pages, splitext(relpath(joinpath(directory, file), dist_dir))[1])
+ end
+ return sort(pages)
+end
+
+# Stubs left at the root by earlier deployments for pages that no longer exist. Only
+# files carrying the marker are candidates, so version folders and hand-written pages are
+# never touched, and `stable` is skipped because its deployment replaces it wholesale.
+function stale_root_stubs(root, current_paths)
+ stale = String[]
+ for (directory, directories, files) in walkdir(root)
+ relative = relpath(directory, root)
+ relative == "." && filter!(!in(PRESERVED_ROOT_DIRS), directories)
+ for file in files
+ endswith(file, ".html") || continue
+ path = relative == "." ? file : joinpath(relative, file)
+ path in current_paths && continue
+ startswith(read(joinpath(root, path), String), STUB_MARKER) &&
+ push!(stale, path)
+ end
+ end
+ return sort(stale)
+end
diff --git a/docs/gen_api.py b/docs/gen_api.py
index 0e9c5a2db..edaa2c041 100644
--- a/docs/gen_api.py
+++ b/docs/gen_api.py
@@ -345,50 +345,12 @@ def process_markdown_file(source_path: Path, output_path: Path) -> None:
print(f"Generated {output_path} from {source_path}")
-def process_readme_to_index(readme_path: Path, output_path: Path) -> None:
- """Convert README.md to index.md with VitePress syntax."""
- content = readme_path.read_text()
-
- # Add the video splash at the top
- video_splash = '\n'
- content = video_splash + content
-
- # Convert blocks to VitePress ::: details syntax
- # Pattern matches: \n\n### Title\n\n
\ncontent\n
- def replace_details(match):
- summary = match.group(1).strip()
- body = match.group(2).strip()
- # Extract title from summary (remove ### prefix if present)
- title = summary.replace("###", "").strip()
- return f"::: details {title}\n\n{body}\n\n:::"
-
- # Match ...
... blocks
- details_pattern = r"\s*\s*(.*?)\s*
\s*(.*?)\s*"
- content = re.sub(details_pattern, replace_details, content, flags=re.DOTALL)
-
- output_path.write_text(content)
- print(f"Generated {output_path} from {readme_path}")
-
-
def main():
"""Main entry point."""
# Get the directory containing this script
script_dir = Path(__file__).parent.resolve()
src_dir = script_dir / "src"
- # Generate index.md from README.md
- readme_path = script_dir.parent / "README.md"
- index_path = src_dir / "index.md"
- if readme_path.exists():
- try:
- process_readme_to_index(readme_path, index_path)
- except Exception as e:
- print(f"Error generating index.md: {e}", file=sys.stderr)
- import traceback
-
- traceback.print_exc()
- sys.exit(1)
-
# Process _api.md → api.md and _api-advanced.md → api-advanced.md
for source_name, output_name in [
("_api.md", "api.md"),
diff --git a/docs/gen_vitepress_index.py b/docs/gen_vitepress_index.py
index 3806d12b0..5ddbab121 100644
--- a/docs/gen_vitepress_index.py
+++ b/docs/gen_vitepress_index.py
@@ -3,7 +3,30 @@
Generate VitePress-compatible index.md by combining _index.md template with README.md content.
"""
+import re
from pathlib import Path
+from urllib.parse import urlsplit
+
+REPO_ROOT = Path(__file__).resolve().parent.parent
+REPO_URL = "https://github.com/astroautomata/PySR"
+
+
+def repo_link(match):
+ url = match.group(2)
+ parts = urlsplit(url)
+ if parts.scheme or parts.netloc or url.startswith(("/", "#")):
+ return match.group(0)
+
+ path = (REPO_ROOT / parts.path).resolve()
+ if not path.is_relative_to(REPO_ROOT) or not path.exists():
+ return match.group(0)
+ if path.is_relative_to(REPO_ROOT / "docs" / "src") and path.suffix == ".md":
+ return match.group(0)
+
+ kind = "tree" if path.is_dir() else "blob"
+ suffix = f"?{parts.query}" if parts.query else ""
+ suffix += f"#{parts.fragment}" if parts.fragment else ""
+ return f"{match.group(1)}{REPO_URL}/{kind}/master/{path.relative_to(REPO_ROOT).as_posix()}{suffix})"
def process_readme_content(readme_content):
@@ -24,20 +47,13 @@ def process_readme_content(readme_content):
if not found_first_header:
continue
- # Skip GitHub user-images URLs
- if "user-images" in line:
- continue
-
# Adjust header levels: # -> ##
if line.startswith("# ") and not line.startswith("# PySR"):
line = "#" + line
- # Transform "### Test status" to "**Test status**"
- if line == "### Test status":
- line = "**Test status**"
-
# Fix internal links to use VitePress paths
- line = line.replace("(https://ai.damtp.cam.ac.uk/pysr/papers)", "(/papers)")
+ line = line.replace("(https://pysr.ai/papers)", "(/papers)")
+ line = re.sub(r"(?=14.6"
}
@@ -1394,6 +1397,7 @@
"version": "4.1.3",
"resolved": "https://registry.npmjs.org/ansi-colors/-/ansi-colors-4.1.3.tgz",
"integrity": "sha512-/6w/C21Pm1A7aZitlI5Ni/2J6FFQN8i1Cvz3kHABAAbw93v/NlvKdVOqz7CCWz/3iv/JplRSEEZ83XION15ovw==",
+ "license": "MIT",
"engines": {
"node": ">=6"
}
@@ -1414,7 +1418,8 @@
"node_modules/boolbase": {
"version": "1.0.0",
"resolved": "https://registry.npmjs.org/boolbase/-/boolbase-1.0.0.tgz",
- "integrity": "sha512-JZOSA7Mo9sNGB8+UjSgzdLtokWAky1zbztM3WRLCbZ70/3cTANmQmOdR7y2g+J0e2WXywy1yS468tY+IruqEww=="
+ "integrity": "sha512-JZOSA7Mo9sNGB8+UjSgzdLtokWAky1zbztM3WRLCbZ70/3cTANmQmOdR7y2g+J0e2WXywy1yS468tY+IruqEww==",
+ "license": "ISC"
},
"node_modules/ccount": {
"version": "2.0.1",
@@ -1447,6 +1452,7 @@
"version": "1.0.0-rc.10",
"resolved": "https://registry.npmjs.org/cheerio/-/cheerio-1.0.0-rc.10.tgz",
"integrity": "sha512-g0J0q/O6mW8z5zxQ3A8E8J1hUgp4SMOvEoW/x84OwyHKe/Zccz83PVT4y5Crcr530FV6NgmKI1qvGTKVl9XXVw==",
+ "license": "MIT",
"dependencies": {
"cheerio-select": "^1.5.0",
"dom-serializer": "^1.3.2",
@@ -1467,6 +1473,7 @@
"version": "1.6.0",
"resolved": "https://registry.npmjs.org/cheerio-select/-/cheerio-select-1.6.0.tgz",
"integrity": "sha512-eq0GdBvxVFbqWgmCm7M3XGs1I8oLy/nExUnh6oLqmBditPO9AqQJrkslDpMun/hZ0yyTs8L0m85OHp4ho6Qm9g==",
+ "license": "BSD-2-Clause",
"dependencies": {
"css-select": "^4.3.0",
"css-what": "^6.0.1",
@@ -1491,6 +1498,7 @@
"version": "6.2.1",
"resolved": "https://registry.npmjs.org/commander/-/commander-6.2.1.tgz",
"integrity": "sha512-U7VdrJFnJgo4xjrHpTzu0yrHPGImdsmD95ZlgYSEajAn2JKzDhDTPG9kBTefmObL2w/ngeZnilk+OV9CG3d7UA==",
+ "license": "MIT",
"engines": {
"node": ">= 6"
}
@@ -1510,6 +1518,7 @@
"version": "4.3.0",
"resolved": "https://registry.npmjs.org/css-select/-/css-select-4.3.0.tgz",
"integrity": "sha512-wPpOYtnsVontu2mODhA19JrqWxNsfdatRKd64kmpRbQgh1KtItko5sTnEpPdpSaJszTOhEMlF/RPz28qj4HqhQ==",
+ "license": "BSD-2-Clause",
"dependencies": {
"boolbase": "^1.0.0",
"css-what": "^6.0.1",
@@ -1525,6 +1534,7 @@
"version": "6.2.2",
"resolved": "https://registry.npmjs.org/css-what/-/css-what-6.2.2.tgz",
"integrity": "sha512-u/O3vwbptzhMs3L1fQE82ZSLHQQfto5gyZzwteVIEyeaY5Fc7R4dapF/BvRoSYFeqfBk4m0V1Vafq5Pjv25wvA==",
+ "license": "BSD-2-Clause",
"engines": {
"node": ">= 6"
},
@@ -1538,9 +1548,10 @@
"integrity": "sha512-M1uQkMl8rQK/szD0LNhtqxIPLpimGm8sOBwU7lLnCpSbTyY3yeU1Vc7l4KT5zT4s/yOxHH5O7tIuuLOCnLADRw=="
},
"node_modules/d3-format": {
- "version": "3.1.0",
- "resolved": "https://registry.npmjs.org/d3-format/-/d3-format-3.1.0.tgz",
- "integrity": "sha512-YyUI6AEuY/Wpt8KWLgZHsIU86atmikuoOmCfommt0LYHiQSPjvX2AcFc38PX0CBpr2RCyZhjex+NS/LPOv6YqA==",
+ "version": "3.1.2",
+ "resolved": "https://registry.npmjs.org/d3-format/-/d3-format-3.1.2.tgz",
+ "integrity": "sha512-AJDdYOdnyRDV5b6ArilzCPPwc1ejkHcoyFarqlPqT7zRYjhavcT3uSrqcMvsgh2CgoPbK3RCwyHaVyxYcP2Arg==",
+ "license": "ISC",
"engines": {
"node": ">=12"
}
@@ -1569,6 +1580,7 @@
"version": "1.4.1",
"resolved": "https://registry.npmjs.org/dom-serializer/-/dom-serializer-1.4.1.tgz",
"integrity": "sha512-VHwB3KfrcOOkelEG2ZOfxqLZdfkil8PtJi4P8N2MMXucZq2yLp75ClViUlOVwyoHEDjYU433Aq+5zWP61+RGag==",
+ "license": "MIT",
"dependencies": {
"domelementtype": "^2.0.1",
"domhandler": "^4.2.0",
@@ -1582,6 +1594,7 @@
"version": "2.2.0",
"resolved": "https://registry.npmjs.org/entities/-/entities-2.2.0.tgz",
"integrity": "sha512-p92if5Nz619I0w+akJrLZH0MX0Pb5DX39XOwQTtXSdQQOaYH03S1uIQp4mhOZtAXrxq4ViO67YTiLBo2638o9A==",
+ "license": "BSD-2-Clause",
"funding": {
"url": "https://github.com/fb55/entities?sponsor=1"
}
@@ -1595,12 +1608,14 @@
"type": "github",
"url": "https://github.com/sponsors/fb55"
}
- ]
+ ],
+ "license": "BSD-2-Clause"
},
"node_modules/domhandler": {
"version": "4.3.1",
"resolved": "https://registry.npmjs.org/domhandler/-/domhandler-4.3.1.tgz",
"integrity": "sha512-GrwoxYN+uWlzO8uhUXRl0P+kHE4GtVPfYzVLcUxPL7KNdHKj66vvlhiweIHqYYXWlw+T8iLMp42Lm67ghw4WMQ==",
+ "license": "BSD-2-Clause",
"dependencies": {
"domelementtype": "^2.2.0"
},
@@ -1615,6 +1630,7 @@
"version": "2.8.0",
"resolved": "https://registry.npmjs.org/domutils/-/domutils-2.8.0.tgz",
"integrity": "sha512-w96Cjofp72M5IIhpjgobBimYEfoPjx1Vx0BSX9P30WBdZW2WIKU0T1Bd0kz2eNZ9ikjKgHbEyKx8BB6H1L3h3A==",
+ "license": "BSD-2-Clause",
"dependencies": {
"dom-serializer": "^1.0.1",
"domelementtype": "^2.2.0",
@@ -1693,6 +1709,7 @@
"version": "3.0.0",
"resolved": "https://registry.npmjs.org/escape-goat/-/escape-goat-3.0.0.tgz",
"integrity": "sha512-w3PwNZJwRxlp47QGzhuEBldEqVHHhh8/tIPcl6ecf2Bou99cdAt0knihBV0Ecc7CGxYduXVBDheH1K2oADRlvw==",
+ "license": "MIT",
"engines": {
"node": ">=10"
},
@@ -1704,6 +1721,7 @@
"version": "3.2.25",
"resolved": "https://registry.npmjs.org/esm/-/esm-3.2.25.tgz",
"integrity": "sha512-U1suiZ2oDVWv4zPO56S0NcR5QriEahGtdN2OR6FiOG4WJvcjBVFB0qI4+eKoWFH483PKGuLuu6V8Z4T5g63UVA==",
+ "license": "MIT",
"engines": {
"node": ">=6"
}
@@ -1799,6 +1817,7 @@
"url": "https://github.com/sponsors/fb55"
}
],
+ "license": "MIT",
"dependencies": {
"domelementtype": "^2.0.1",
"domhandler": "^4.0.0",
@@ -1810,6 +1829,7 @@
"version": "2.2.0",
"resolved": "https://registry.npmjs.org/entities/-/entities-2.2.0.tgz",
"integrity": "sha512-p92if5Nz619I0w+akJrLZH0MX0Pb5DX39XOwQTtXSdQQOaYH03S1uIQp4mhOZtAXrxq4ViO67YTiLBo2638o9A==",
+ "license": "BSD-2-Clause",
"funding": {
"url": "https://github.com/fb55/entities?sponsor=1"
}
@@ -1847,6 +1867,7 @@
"version": "8.1.0",
"resolved": "https://registry.npmjs.org/juice/-/juice-8.1.0.tgz",
"integrity": "sha512-FLzurJrx5Iv1e7CfBSZH68dC04EEvXvvVvPYB7Vx1WAuhCp1ZPIMtqxc+WTWxVkpTIC2Ach/GAv0rQbtGf6YMA==",
+ "license": "MIT",
"dependencies": {
"cheerio": "1.0.0-rc.10",
"commander": "^6.1.0",
@@ -1887,9 +1908,20 @@
}
},
"node_modules/linkify-it": {
- "version": "5.0.0",
- "resolved": "https://registry.npmjs.org/linkify-it/-/linkify-it-5.0.0.tgz",
- "integrity": "sha512-5aHCbzQRADcdP+ATqnDuhhJ/MRIqDkZX5pyjFHRRysS8vZ5AbqGEoFIb6pYHPZ+L/OC2Lc+xT8uHVVR5CAK/wQ==",
+ "version": "5.0.2",
+ "resolved": "https://registry.npmjs.org/linkify-it/-/linkify-it-5.0.2.tgz",
+ "integrity": "sha512-ONTm2jCMAVZjgQa/Fy1kScXsuOoF5NPTsoFBdE1KVIZ2vAh/r9+Bqo+0jINCBYnavTPQZz38QzFTme79ENoN3Q==",
+ "funding": [
+ {
+ "type": "github",
+ "url": "https://github.com/sponsors/puzrin"
+ },
+ {
+ "type": "github",
+ "url": "https://github.com/sponsors/markdown-it"
+ }
+ ],
+ "license": "MIT",
"dependencies": {
"uc.micro": "^2.0.0"
}
@@ -1921,13 +1953,24 @@
"integrity": "sha512-1I+1qpDt4idfgLQG+BNWmrqku+7/2bi5nLf4YwF8y8zXvmfiTBY3PV3ZibfrjBueCByROpuBjLLFCajqkgYoLQ=="
},
"node_modules/markdown-it": {
- "version": "14.1.0",
- "resolved": "https://registry.npmjs.org/markdown-it/-/markdown-it-14.1.0.tgz",
- "integrity": "sha512-a54IwgWPaeBCAAsv13YgmALOF1elABB08FxO9i+r4VFk5Vl4pKokRPeX8u5TCgSsPi6ec1otfLjdOpVcgbpshg==",
+ "version": "14.3.0",
+ "resolved": "https://registry.npmjs.org/markdown-it/-/markdown-it-14.3.0.tgz",
+ "integrity": "sha512-RCEsPjR+sr0x+AuYp601tKTkgFG4YEPLCzHST3cQ/fhlJkqAkz1L2/Qbp1j9qw5SBwQHFBoW8+hoN5xssOF0Tw==",
+ "funding": [
+ {
+ "type": "github",
+ "url": "https://github.com/sponsors/puzrin"
+ },
+ {
+ "type": "github",
+ "url": "https://github.com/sponsors/markdown-it"
+ }
+ ],
+ "license": "MIT",
"dependencies": {
"argparse": "^2.0.1",
- "entities": "^4.4.0",
- "linkify-it": "^5.0.0",
+ "entities": "^4.5.0",
+ "linkify-it": "^5.0.2",
"mdurl": "^2.0.0",
"punycode.js": "^2.3.1",
"uc.micro": "^2.1.0"
@@ -1945,15 +1988,17 @@
"version": "4.3.2",
"resolved": "https://registry.npmjs.org/markdown-it-mathjax3/-/markdown-it-mathjax3-4.3.2.tgz",
"integrity": "sha512-TX3GW5NjmupgFtMJGRauioMbbkGsOXAAt1DZ/rzzYmTHqzkO1rNAdiMD4NiruurToPApn2kYy76x02QN26qr2w==",
+ "license": "MIT",
"dependencies": {
"juice": "^8.0.0",
"mathjax-full": "^3.2.0"
}
},
"node_modules/mathjax-full": {
- "version": "3.2.2",
- "resolved": "https://registry.npmjs.org/mathjax-full/-/mathjax-full-3.2.2.tgz",
- "integrity": "sha512-+LfG9Fik+OuI8SLwsiR02IVdjcnRCy5MufYLi0C3TdMT56L/pjB0alMVGgoWJF8pN9Rc7FESycZB9BMNWIid5w==",
+ "version": "3.2.1",
+ "resolved": "https://registry.npmjs.org/mathjax-full/-/mathjax-full-3.2.1.tgz",
+ "integrity": "sha512-aUz9o16MGZdeiIBwZjAfUBTiJb7LRqzZEl1YOZ8zQMGYIyh1/nxRebxKxjDe9L+xcZCr2OHdzoFBMcd6VnLv9Q==",
+ "license": "Apache-2.0",
"dependencies": {
"esm": "^3.2.25",
"mhchemparser": "^4.1.0",
@@ -1989,12 +2034,14 @@
"node_modules/mensch": {
"version": "0.3.4",
"resolved": "https://registry.npmjs.org/mensch/-/mensch-0.3.4.tgz",
- "integrity": "sha512-IAeFvcOnV9V0Yk+bFhYR07O3yNina9ANIN5MoXBKYJ/RLYPurd2d0yw14MDhpr9/momp0WofT1bPUh3hkzdi/g=="
+ "integrity": "sha512-IAeFvcOnV9V0Yk+bFhYR07O3yNina9ANIN5MoXBKYJ/RLYPurd2d0yw14MDhpr9/momp0WofT1bPUh3hkzdi/g==",
+ "license": "MIT"
},
"node_modules/mhchemparser": {
"version": "4.2.1",
"resolved": "https://registry.npmjs.org/mhchemparser/-/mhchemparser-4.2.1.tgz",
- "integrity": "sha512-kYmyrCirqJf3zZ9t/0wGgRZ4/ZJw//VwaRVGA75C4nhE60vtnIzhl9J9ndkX/h6hxSN7pjg/cE0VxbnNM+bnDQ=="
+ "integrity": "sha512-kYmyrCirqJf3zZ9t/0wGgRZ4/ZJw//VwaRVGA75C4nhE60vtnIzhl9J9ndkX/h6hxSN7pjg/cE0VxbnNM+bnDQ==",
+ "license": "Apache-2.0"
},
"node_modules/micromark-util-character": {
"version": "2.1.1",
@@ -2105,7 +2152,8 @@
"node_modules/mj-context-menu": {
"version": "0.6.1",
"resolved": "https://registry.npmjs.org/mj-context-menu/-/mj-context-menu-0.6.1.tgz",
- "integrity": "sha512-7NO5s6n10TIV96d4g2uDpG7ZDpIhMh0QNfGdJw/W47JswFcosz457wqz/b5sAKvl12sxINGFCn80NZHKwxQEXA=="
+ "integrity": "sha512-7NO5s6n10TIV96d4g2uDpG7ZDpIhMh0QNfGdJw/W47JswFcosz457wqz/b5sAKvl12sxINGFCn80NZHKwxQEXA==",
+ "license": "Apache-2.0"
},
"node_modules/nanoid": {
"version": "3.3.11",
@@ -2144,6 +2192,7 @@
"version": "2.7.0",
"resolved": "https://registry.npmjs.org/node-fetch/-/node-fetch-2.7.0.tgz",
"integrity": "sha512-c4FRfUm/dbcWZ7U+1Wq0AwCyFL+3nt2bEw05wfxSz+DWpWsitgmSgYmy2dQdWyKC1694ELPqMs/YzUSNozLt8A==",
+ "license": "MIT",
"dependencies": {
"whatwg-url": "^5.0.0"
},
@@ -2163,6 +2212,7 @@
"version": "2.1.1",
"resolved": "https://registry.npmjs.org/nth-check/-/nth-check-2.1.1.tgz",
"integrity": "sha512-lqjrjmaOoAnWfMmBPL+XNnynZh2+swxiX3WUE0s4yEHI6m+AwrK2UZOimIRl3X/4QctVqS8AiZjFqyOGrMXb/w==",
+ "license": "BSD-2-Clause",
"dependencies": {
"boolbase": "^1.0.0"
},
@@ -2191,12 +2241,14 @@
"node_modules/parse5": {
"version": "6.0.1",
"resolved": "https://registry.npmjs.org/parse5/-/parse5-6.0.1.tgz",
- "integrity": "sha512-Ofn/CTFzRGTTxwpNEs9PP93gXShHcTq255nzRYSKe8AkVpZY7e1fpmTfOyoIvjP5HG7Z2ZM7VS9PPhQGW2pOpw=="
+ "integrity": "sha512-Ofn/CTFzRGTTxwpNEs9PP93gXShHcTq255nzRYSKe8AkVpZY7e1fpmTfOyoIvjP5HG7Z2ZM7VS9PPhQGW2pOpw==",
+ "license": "MIT"
},
"node_modules/parse5-htmlparser2-tree-adapter": {
"version": "6.0.1",
"resolved": "https://registry.npmjs.org/parse5-htmlparser2-tree-adapter/-/parse5-htmlparser2-tree-adapter-6.0.1.tgz",
"integrity": "sha512-qPuWvbLgvDGilKc5BoicRovlT4MtYT6JfJyBOMDsKoiT+GiuP5qyrPCnR9HcPECIJJmZh5jRndyNThnhhb/vlA==",
+ "license": "MIT",
"dependencies": {
"parse5": "^6.0.1"
}
@@ -2391,6 +2443,7 @@
"version": "1.12.2",
"resolved": "https://registry.npmjs.org/slick/-/slick-1.12.2.tgz",
"integrity": "sha512-4qdtOGcBjral6YIBCWJ0ljFSKNLz9KkhbWtuGvUyRowl1kxfuE1x/Z/aJcaiilpb3do9bl5K7/1h9XC5wWpY/A==",
+ "license": "MIT (http://mootools.net/license.txt)",
"engines": {
"node": "*"
}
@@ -2430,11 +2483,12 @@
}
},
"node_modules/speech-rule-engine": {
- "version": "4.1.2",
- "resolved": "https://registry.npmjs.org/speech-rule-engine/-/speech-rule-engine-4.1.2.tgz",
- "integrity": "sha512-S6ji+flMEga+1QU79NDbwZ8Ivf0S/MpupQQiIC0rTpU/ZTKgcajijJJb1OcByBQDjrXCN1/DJtGz4ZJeBMPGJw==",
+ "version": "4.1.4",
+ "resolved": "https://registry.npmjs.org/speech-rule-engine/-/speech-rule-engine-4.1.4.tgz",
+ "integrity": "sha512-i/VCLG1fvRc95pMHRqG4aQNscv+9aIsqA2oI7ZQS51sTdUcDHYX6cpT8/tqZ+enjs1tKVwbRBWgxut9SWn+f9g==",
+ "license": "Apache-2.0",
"dependencies": {
- "@xmldom/xmldom": "0.9.8",
+ "@xmldom/xmldom": "0.9.10",
"commander": "13.1.0",
"wicked-good-xpath": "1.3.0"
},
@@ -2446,6 +2500,7 @@
"version": "13.1.0",
"resolved": "https://registry.npmjs.org/commander/-/commander-13.1.0.tgz",
"integrity": "sha512-/rFeCpNJQbhSZjGVwO9RFV3xPqbnERS8MmIQzCtD/zl6gpJuV/bMLuN92oG3F7d8oDEHHRrujSXNUr8fpjntKw==",
+ "license": "MIT",
"engines": {
"node": ">=18"
}
@@ -2507,7 +2562,8 @@
"node_modules/tr46": {
"version": "0.0.3",
"resolved": "https://registry.npmjs.org/tr46/-/tr46-0.0.3.tgz",
- "integrity": "sha512-N3WMsuqV66lT30CrXNbEjx4GEwlow3v6rr4mCcv6prnfwhS01rkgyFdjPNBYd9br7LpXV1+Emh01fHnq2Gdgrw=="
+ "integrity": "sha512-N3WMsuqV66lT30CrXNbEjx4GEwlow3v6rr4mCcv6prnfwhS01rkgyFdjPNBYd9br7LpXV1+Emh01fHnq2Gdgrw==",
+ "license": "MIT"
},
"node_modules/trim-lines": {
"version": "3.0.1",
@@ -2526,7 +2582,8 @@
"node_modules/uc.micro": {
"version": "2.1.0",
"resolved": "https://registry.npmjs.org/uc.micro/-/uc.micro-2.1.0.tgz",
- "integrity": "sha512-ARDJmphmdvUk6Glw7y9DQ2bFkKBHwQHLi2lsaH6PPmz/Ka9sFOBsBluozhDltWmnv9u/cF6Rt87znRTPV+yp/A=="
+ "integrity": "sha512-ARDJmphmdvUk6Glw7y9DQ2bFkKBHwQHLi2lsaH6PPmz/Ka9sFOBsBluozhDltWmnv9u/cF6Rt87znRTPV+yp/A==",
+ "license": "MIT"
},
"node_modules/unist-util-is": {
"version": "6.0.0",
@@ -2595,6 +2652,7 @@
"version": "3.0.1",
"resolved": "https://registry.npmjs.org/valid-data-url/-/valid-data-url-3.0.1.tgz",
"integrity": "sha512-jOWVmzVceKlVVdwjNSenT4PbGghU0SBIizAev8ofZVgivk/TVHXSbNL8LP6M3spZvkR9/QolkyJavGSX5Cs0UA==",
+ "license": "MIT",
"engines": {
"node": ">=10"
}
@@ -2724,12 +2782,17 @@
}
},
"node_modules/vitepress-plugin-tabs": {
- "version": "0.6.0",
- "resolved": "https://registry.npmjs.org/vitepress-plugin-tabs/-/vitepress-plugin-tabs-0.6.0.tgz",
- "integrity": "sha512-35mQUbI6lcVrLj0PpwautXm8LiV6Gz5ZDlklVRf/0KDd6gGs8HaZztqOxdffnzCiKIFLURBx930xq28it9jQTg==",
+ "version": "0.9.0",
+ "resolved": "https://registry.npmjs.org/vitepress-plugin-tabs/-/vitepress-plugin-tabs-0.9.0.tgz",
+ "integrity": "sha512-OB+/T6SkwGIQHIQi1NLl011kWCMzfsVEx0/q+NfF6EmvvpcfDzRClc/iEwUlq8TkQnFxANu+GTdmvtXEjJM2Lg==",
+ "license": "MIT",
+ "funding": {
+ "type": "github",
+ "url": "https://github.com/sponsors/sapphi-red"
+ },
"peerDependencies": {
- "vitepress": "^1.0.0",
- "vue": "^3.3.8"
+ "vitepress": "^1.0.0 || ^2.0.0-alpha.17",
+ "vue": "^3.5.0"
}
},
"node_modules/vue": {
@@ -2756,6 +2819,7 @@
"version": "6.0.1",
"resolved": "https://registry.npmjs.org/web-resource-inliner/-/web-resource-inliner-6.0.1.tgz",
"integrity": "sha512-kfqDxt5dTB1JhqsCUQVFDj0rmY+4HLwGQIsLPbyrsN9y9WV/1oFDSx3BQ4GfCv9X+jVeQ7rouTqwK53rA/7t8A==",
+ "license": "MIT",
"dependencies": {
"ansi-colors": "^4.1.1",
"escape-goat": "^3.0.0",
@@ -2772,6 +2836,7 @@
"version": "3.3.0",
"resolved": "https://registry.npmjs.org/domhandler/-/domhandler-3.3.0.tgz",
"integrity": "sha512-J1C5rIANUbuYK+FuFL98650rihynUOEzRLxW+90bKZRWB6A1X1Tf82GxR1qAWLyfNPRvjqfip3Q5tdYlmAa9lA==",
+ "license": "BSD-2-Clause",
"dependencies": {
"domelementtype": "^2.0.1"
},
@@ -2786,6 +2851,7 @@
"version": "2.2.0",
"resolved": "https://registry.npmjs.org/entities/-/entities-2.2.0.tgz",
"integrity": "sha512-p92if5Nz619I0w+akJrLZH0MX0Pb5DX39XOwQTtXSdQQOaYH03S1uIQp4mhOZtAXrxq4ViO67YTiLBo2638o9A==",
+ "license": "BSD-2-Clause",
"funding": {
"url": "https://github.com/fb55/entities?sponsor=1"
}
@@ -2794,6 +2860,7 @@
"version": "5.0.1",
"resolved": "https://registry.npmjs.org/htmlparser2/-/htmlparser2-5.0.1.tgz",
"integrity": "sha512-vKZZra6CSe9qsJzh0BjBGXo8dvzNsq/oGvsjfRdOrrryfeD9UOBEEQdeoqCRmKZchF5h2zOBMQ6YuQ0uRUmdbQ==",
+ "license": "MIT",
"dependencies": {
"domelementtype": "^2.0.1",
"domhandler": "^3.3.0",
@@ -2808,6 +2875,7 @@
"version": "2.6.0",
"resolved": "https://registry.npmjs.org/mime/-/mime-2.6.0.tgz",
"integrity": "sha512-USPkMeET31rOMiarsBNIHZKLGgvKc/LrjofAnBlOttf5ajRvqiRA8QsenbcooctK6d6Ts6aqZXBA+XbkKthiQg==",
+ "license": "MIT",
"bin": {
"mime": "cli.js"
},
@@ -2818,12 +2886,14 @@
"node_modules/webidl-conversions": {
"version": "3.0.1",
"resolved": "https://registry.npmjs.org/webidl-conversions/-/webidl-conversions-3.0.1.tgz",
- "integrity": "sha512-2JAn3z8AR6rjK8Sm8orRC0h/bcl/DqL7tRPdGZ4I1CjdF+EaMLmYxBHyXuKL849eucPFhvBoxMsflfOb8kxaeQ=="
+ "integrity": "sha512-2JAn3z8AR6rjK8Sm8orRC0h/bcl/DqL7tRPdGZ4I1CjdF+EaMLmYxBHyXuKL849eucPFhvBoxMsflfOb8kxaeQ==",
+ "license": "BSD-2-Clause"
},
"node_modules/whatwg-url": {
"version": "5.0.0",
"resolved": "https://registry.npmjs.org/whatwg-url/-/whatwg-url-5.0.0.tgz",
"integrity": "sha512-saE57nupxk6v3HY35+jzBwYa0rKSy0XR8JSxZPwgLr7ys0IBzhGviA1/TUGJLmSVqs8pb9AnvICXEuOHLprYTw==",
+ "license": "MIT",
"dependencies": {
"tr46": "~0.0.3",
"webidl-conversions": "^3.0.0"
@@ -2832,7 +2902,8 @@
"node_modules/wicked-good-xpath": {
"version": "1.3.0",
"resolved": "https://registry.npmjs.org/wicked-good-xpath/-/wicked-good-xpath-1.3.0.tgz",
- "integrity": "sha512-Gd9+TUn5nXdwj/hFsPVx5cuHHiF5Bwuc30jZ4+ronF1qHK5O7HD0sgmXWSEgwKquT3ClLoKPVbO6qGwVwLzvAw=="
+ "integrity": "sha512-Gd9+TUn5nXdwj/hFsPVx5cuHHiF5Bwuc30jZ4+ronF1qHK5O7HD0sgmXWSEgwKquT3ClLoKPVbO6qGwVwLzvAw==",
+ "license": "MIT"
},
"node_modules/zwitch": {
"version": "2.0.4",
diff --git a/docs/package.json b/docs/package.json
index dfec562ea..e92aaa1df 100644
--- a/docs/package.json
+++ b/docs/package.json
@@ -11,11 +11,11 @@
},
"dependencies": {
"@nolebase/vitepress-plugin-enhanced-readabilities": "^2.14.0",
- "d3-format": "^3.1.0",
- "markdown-it": "^14.1.0",
+ "d3-format": "^3.1.2",
+ "markdown-it": "^14.2.0",
"markdown-it-footnote": "^4.0.0",
"markdown-it-mathjax3": "^4.3.2",
"vitepress": "^1.6.3",
- "vitepress-plugin-tabs": "^0.6.0"
+ "vitepress-plugin-tabs": "^0.9.0"
}
}
diff --git a/docs/papers.yml b/docs/papers.yml
index 21afd44fd..2709f1152 100644
--- a/docs/papers.yml
+++ b/docs/papers.yml
@@ -2,6 +2,35 @@
# information to generate the "Research Showcase"
papers:
+ - title: "Maximum updraft velocity beyond CAPE: the role of boundary layer dynamics and pressure perturbations"
+ authors:
+ - Bety Pechacova (1)
+ - Alejandro Casallas (1)
+ - Tom Beucler (2,3)
+ - Lokahith Agasthya (4)
+ - Caroline Muller (1)
+ affiliations:
+ 1: Institute of Science and Technology Austria
+ 2: Faculty of Geosciences and Environment, University of Lausanne
+ 3: Expertise Center for Climate Extremes, University of Lausanne
+ 4: Department of Meteorology and Geophysics, University of Vienna
+ link: https://doi.org/10.48550/arXiv.2607.23627
+ abstract: "Deep convective updraft velocities play a key role in the Earth's climate system, influencing precipitation extremes, lightning, and the planetary energy budget. While Convective Available Potential Energy (CAPE) is widely used to explain maximum updraft velocity ($w_{\\max}$), CAPE is an imperfect predictor as updrafts are also influenced by entrainment, boundary layer dynamics, pressure perturbations, and condensate loading. However, the relative importance of these processes and how they interact to set $w_{\\max}$ in individual clouds remains unclear. Here, we use equation learning to identify compact, physically interpretable relationships linking environmental and in-cloud conditions to $w_{\\max}$ in individual tracked clouds across idealized radiative–convective equilibrium regimes spanning a range of sea surface temperatures and radiative cooling rates. For pre-storm prediction, CAPE and local mean boundary layer vertical velocity ($\\overline{w_{\\mathrm{bl}}}$) together explain nearly half the variance in $w_{\\max}$ across regimes ($R^2=0.47$). While CAPE captures regime-mean differences, it has little predictive value within a single simulation. $\\overline{w_{\\mathrm{bl}}}$ is essential for capturing cloud-to-cloud variability, including the suppression of $w_{\\max}$ even at high $\\mathrm{CAPE}$ values. At the time of peak intensity, a simple approximate Bernoulli-like invariant combining maximum pressure perturbation and maximum cloud condensate explains 89% of the variance ($R^2=0.89$). The tight link between $w_{\\max}$ and pressure perturbation supports the sticky thermals hypothesis and highlights the importance of dynamic pressure effects, often neglected in updraft theories. These results highlight $\\overline{w_{\\mathrm{bl}}}$ as an important regulator of convective intensity alongside CAPE, and demonstrate that dynamic pressure plays an important role within individual updrafts."
+ image: https://raw.githubusercontent.com/MilesCranmer/PySR_Docs/d6a8df0a069efb549ed9fc53e87cf0877ad3fee6/images/pr1295_updraft_velocity.png
+ date: 2026-07-26
+ - title: Improvement of spatiotemporal generalization in radiocesium transfer models for wheat using symbolic regression
+ authors:
+ - Katsutoshi Seki (1)
+ - Noriko Yamaguchi (2)
+ - Tetsuya Eguchi (2)
+ - Masato Igura (2)
+ affiliations:
+ 1: Natural Science Laboratory, Toyo University
+ 2: Institute for Agro-Environmental Sciences, NARO
+ link: https://doi.org/10.1016/j.jenvrad.2026.108077
+ abstract: "Accurate prediction of the soil-to-plant transfer factor (TF) of radiocesium ($^{137}$Cs) is essential for assessing radionuclide transfer to the human food chain and supporting radiological risk management after nuclear contamination. Semi-empirical TF models based on exchangeable potassium (K) and soil fixation capacity are widely used, but their predictive performance often deteriorates when applied to unobserved sites or years. In this study, we evaluated whether symbolic regression (SR), an approach that derives interpretable mathematical equations from data, can improve the spatiotemporal generalization of TF models. We analyzed long-term monitoring data from 11 upland wheat fields in Japan collected over five survey years ($n$ = 36). Model performance was assessed using structured cross-validation, including leave-one-site-out (LOSO) and leave-one-year-out (LOYO), to explicitly test spatial and temporal extrapolation. A simple K-based model substantially improved baseline predictions compared with a conventional semi-empirical model. However, linear additive extensions incorporating radiocesium interception potential (RIP), cation exchange capacity, pH, and soil $^{137}$Cs did not consistently improve predictive performance under LOSO or LOYO. In contrast, the selected SR model showed the lowest point-estimate prediction errors in both spatial and temporal validations. The selected equation represented a non-linear interaction between exchangeable K and RIP, consistent with reduced K sensitivity in soils with strong specific fixation sites. However, the selected SR-derived equation remains empirical and is valid only within the observed soil-property range and may show nonphysical behavior under extreme RIP conditions. These results suggest that interpretable non-linear equations derived by SR may improve TF prediction within a restricted soil domain."
+ image: https://raw.githubusercontent.com/MilesCranmer/PySR_Docs/d6a8df0a069efb549ed9fc53e87cf0877ad3fee6/images/pr1206_radiocesium_transfer_models.png
+ date: 2026-06-13
- title: Discovering data-driven microbial growth models with symbolic regression
authors:
- T. Anthony Sun (1)
@@ -80,7 +109,7 @@ papers:
- Pericles Panagiotou (1)
affiliations:
1: Laboratory of Fluid Mechanics and Turbomachinery, Department of Mechanical Engineering, Aristotle University of Thessaloniki, 54124 Thessaloniki, Greece
- link: https://www.aifluids.net/proceedings/S6P16.pdf
+ link: https://www.aifluids.net/proceedings/C.%20Pliakos,%20G.%20Efrem,%20D.%20Terzis,%20P.%20Panagiotou,%20Data-driven%20skin%20friction%20estimation%20for%20UAV%20wings%20in%20subsonic%20flows.pdf
abstract: Accurate estimation of the skin friction coefficient (𝐶𝑓) is essential for estimating the wall shear stresses (𝜏𝑤) and ultimately the first-layer cell height (𝑦) in wall-resolved RANS simulations of wings, where turbulence models are used, demanding a specific grid resolution near walls (primarily the 𝑦𝑡𝑎𝑟𝑔𝑒𝑡⁺). Conventional flat-plate correlations often fail to account for the three-dimensional nature of real wing flows, introducing uncertainties in 𝐶𝑓 predictions and leading to multiple CFD analyses and mesh refinements to meet the targets. In this work, we propose a machine-learning-based approach exploring symbolic regression to derive a model that correlates wing-specific parameters (e.g., Reynolds number, angle of attack, thickness-to-chord ratio, wing sweep angle) with 𝐶𝑓 at the Mean Aerodynamic Chord (MAC). Data are acquired from an in-house database of over 5,000 RANS simulations for UAV wings operating in the low subsonic regime, covering a wide design space, all conducted following best-practice CFD guidelines to ensure high fidelity. These analyses are performed at various flow conditions covering Reynolds numbers from 10⁵ to 10⁷ and include the complete drag polar for each wing. The proposed correlation provides improved agreement with CFD data and enables more accurate 𝑦⁺ estimations. Validation on different wing geometries, including the ONERA M6 and in-house UAV wings, confirmed the robustness of the model, which improves boundary-layer resolution with only a marginal (~2%) increase in total mesh size, while achieving an R² of 0.68 with negligible computational inference cost. This explicit, data-driven equation offers an efficient method for streamlining mesh generation in aerodynamic simulations.
image: https://raw.githubusercontent.com/MilesCranmer/PySR_Docs/48020babcca1e73d2f5bbbd4ef3f0ec3265c59fb/images/pr1127_skin_friction.png
date: 2025-05-30
diff --git a/docs/src/.vitepress/config.mts b/docs/src/.vitepress/config.mts
index 1921fdd00..7572d3fda 100644
--- a/docs/src/.vitepress/config.mts
+++ b/docs/src/.vitepress/config.mts
@@ -1,20 +1,31 @@
+import { readFileSync } from 'node:fs'
+import path from 'node:path'
+import { fileURLToPath } from 'node:url'
import { defineConfig } from 'vitepress'
import { tabsMarkdownPlugin } from 'vitepress-plugin-tabs'
import mathjax3 from "markdown-it-mathjax3";
import footnote from "markdown-it-footnote";
-function getBaseRepository(base: string): string {
- if (!base || base === '/') return '/';
- const parts = base.split('/').filter(Boolean);
- return parts.length > 0 ? `/${parts[0]}/` : '/';
+// Maps the juliapkg.json pin to the SymbolicRegression.jl docs folder covering
+// every version the pin admits, following Julia Pkg compat semantics:
+// `=X.Y.Z` is exact, `~X.Y.Z` allows patch bumps, `^X.Y.Z` (or bare) allows minor bumps.
+export function symbolicRegressionDocsSubfolder(pkg: { version?: unknown; rev?: unknown }): string {
+ const compat = typeof pkg.version === 'string' ? pkg.version.trim().match(/^([=~^]?)(\d+)(?:\.(\d+))?(?:\.(\d+))?$/) : null
+ if (compat) {
+ const [, op, major, minor, patch] = compat
+ if (op === '=' && patch !== undefined) return `v${major}.${minor}.${patch}`
+ if ((op === '~' || op === '=' || major === '0') && minor !== undefined) return `v${major}.${minor}`
+ return `v${major}`
+ }
+ return typeof pkg.rev === 'string' && /^v\d+\.\d+\.\d+$/.test(pkg.rev) ? pkg.rev : 'dev'
}
-// Canonical domain for SEO - will be replaced by deploy.jl at build time
-// Empty string means no canonical tag (this site is canonical)
-const canonicalDomain = '';
+const configDirectory = path.dirname(fileURLToPath(import.meta.url))
+const juliaPkg = JSON.parse(readFileSync(path.resolve(configDirectory, '../../../pysr/juliapkg.json'), 'utf8'))
+const symbolicRegressionDocsVersion = symbolicRegressionDocsSubfolder(juliaPkg.packages.SymbolicRegression)
const baseTemp = {
- base: '/pysr/',
+ base: '/',
}
const nav = [
@@ -23,13 +34,13 @@ const nav = [
{ text: 'API', link: '/api' },
{
text: 'Forum',
- link: 'https://github.com/MilesCranmer/PySR/discussions'
+ link: 'https://github.com/astroautomata/PySR/discussions'
},
{
text: 'Python',
items: [
{ text: 'Python', link: '/' },
- { text: 'Julia', link: 'https://ai.damtp.cam.ac.uk/symbolicregression/dev/' }
+ { text: 'Julia', link: `https://julia.pysr.ai/${symbolicRegressionDocsVersion}/`, target: '_self', noIcon: true }
]
},
{
@@ -39,7 +50,7 @@ const nav = [
// https://vitepress.dev/reference/site-config
export default defineConfig({
- base: '/pysr/',
+ base: '/',
title: 'PySR',
description: 'High-Performance Symbolic Regression in Python and Julia',
lastUpdated: true,
@@ -47,25 +58,17 @@ export default defineConfig({
outDir: '../dist',
srcExclude: ['**/_*.md'],
head: [
- ['link', { rel: 'icon', href: `${baseTemp.base}favicon.png` }],
- ['script', {src: `${getBaseRepository(baseTemp.base)}versions.js`}],
+ ['link', { rel: 'icon', type: 'image/x-icon', href: `${baseTemp.base}favicon.ico` }],
+ ['link', { rel: 'icon', type: 'image/png', sizes: '16x16', href: `${baseTemp.base}favicon-16x16.png` }],
+ ['link', { rel: 'icon', type: 'image/png', sizes: '32x32', href: `${baseTemp.base}favicon-32x32.png` }],
+ ['link', { rel: 'apple-touch-icon', sizes: '180x180', href: `${baseTemp.base}apple-touch-icon.png` }],
+ ['script', {src: '/versions.js'}],
['script', {src: `${baseTemp.base}siteinfo.js`}]
],
ignoreDeadLinks: true,
- transformHead: ({ pageData }) => {
- if (!canonicalDomain) return [];
-
- // Construct canonical URL for this page
- const pagePath = pageData.relativePath.replace(/((^|\/)index)?\.md$/, '$2');
- const canonicalUrl = `${canonicalDomain}${pagePath}`;
-
- return [
- ['link', { rel: 'canonical', href: canonicalUrl }]
- ];
- },
vite: {
define: {
- __DEPLOY_ABSPATH__: JSON.stringify(getBaseRepository(baseTemp.base)),
+ __DEPLOY_ABSPATH__: JSON.stringify('/'),
},
optimizeDeps: {
exclude: [
@@ -96,7 +99,7 @@ export default defineConfig({
},
themeConfig: {
outline: 'deep',
- logo: 'https://ai.damtp.cam.ac.uk/symbolicregression/dev/logo.png',
+ logo: 'https://raw.githubusercontent.com/MilesCranmer/PySR_Docs/master/images/pysr-logo.png',
search: {
provider: 'local',
options: {
@@ -109,7 +112,21 @@ export default defineConfig({
text: 'Getting Started',
items: [
{ text: 'Introduction', link: '/' },
- { text: 'Examples', link: '/examples' },
+ {
+ text: 'Examples',
+ link: '/examples',
+ collapsed: true,
+ items: [
+ { text: 'Getting started', link: '/examples/getting-started' },
+ { text: 'Expression specifications', link: '/examples/expression-specifications' },
+ { text: 'Objectives and losses', link: '/examples/objectives' },
+ { text: 'Physics and units', link: '/examples/physics' },
+ { text: 'Search behaviour', link: '/examples/search-behaviour' },
+ { text: 'Instrumentation and workflow', link: '/examples/instrumentation' },
+ { text: 'Value types', link: '/examples/value-types' },
+ { text: 'Beyond numeric values', link: '/examples/beyond-numeric-values' },
+ ]
+ },
]
},
{
@@ -131,16 +148,18 @@ export default defineConfig({
text: 'Advanced',
items: [
{ text: 'Tuning', link: '/tuning' },
+ { text: 'Slurm', link: '/slurm' },
{ text: 'Backend', link: '/backend' },
+ { text: 'Migrating from v1', link: '/migration' },
]
}
],
editLink: {
- pattern: 'https://github.com/MilesCranmer/PySR/edit/master/docs/:path',
+ pattern: 'https://github.com/astroautomata/PySR/edit/master/docs/:path',
text: 'Edit this page on GitHub'
},
socialLinks: [
- { icon: 'github', link: 'https://github.com/MilesCranmer/PySR' }
+ { icon: 'github', link: 'https://github.com/astroautomata/PySR' }
],
footer: {
message: 'Made with VitePress',
diff --git a/docs/src/.vitepress/theme/style.css b/docs/src/.vitepress/theme/style.css
index e5b79d3d1..d9afca052 100644
--- a/docs/src/.vitepress/theme/style.css
+++ b/docs/src/.vitepress/theme/style.css
@@ -108,3 +108,14 @@ figure img {
.VPNavBarMenuGroup:hover button .text {
color: var(--vp-c-text-2) !important;
}
+
+/* A section whose holds a heading: keep them on one line */
+.vp-doc summary:has(> h3) {
+ cursor: pointer;
+ margin-top: 32px;
+}
+
+.vp-doc summary > h3 {
+ display: inline;
+ margin: 0;
+}
diff --git a/docs/src/_api_template.md b/docs/src/_api_template.md
index 28cfa3155..2e25d7e86 100644
--- a/docs/src/_api_template.md
+++ b/docs/src/_api_template.md
@@ -75,13 +75,15 @@ PARAMSKEY
heading_level: 3
show_root_full_path: false
-::: pysr.ParametricExpressionSpec
+::: pysr.AbstractExpressionSpec
options:
show_root_heading: true
heading_level: 3
show_root_full_path: false
-::: pysr.AbstractExpressionSpec
+## Type Specifications
+
+::: pysr.TypeSpec
options:
show_root_heading: true
heading_level: 3
@@ -100,3 +102,121 @@ PARAMSKEY
show_root_heading: true
heading_level: 3
show_root_full_path: false
+
+## Mutations
+
+::: pysr.ConstantMutation
+ options:
+ show_root_heading: true
+ heading_level: 3
+ show_root_full_path: false
+
+::: pysr.OperatorMutation
+ options:
+ show_root_heading: true
+ heading_level: 3
+ show_root_full_path: false
+
+::: pysr.FeatureMutation
+ options:
+ show_root_heading: true
+ heading_level: 3
+ show_root_full_path: false
+
+::: pysr.SwapOperandsMutation
+ options:
+ show_root_heading: true
+ heading_level: 3
+ show_root_full_path: false
+
+::: pysr.AddNodeMutation
+ options:
+ show_root_heading: true
+ heading_level: 3
+ show_root_full_path: false
+
+::: pysr.InsertNodeMutation
+ options:
+ show_root_heading: true
+ heading_level: 3
+ show_root_full_path: false
+
+::: pysr.DeleteNodeMutation
+ options:
+ show_root_heading: true
+ heading_level: 3
+ show_root_full_path: false
+
+::: pysr.RotateTreeMutation
+ options:
+ show_root_heading: true
+ heading_level: 3
+ show_root_full_path: false
+
+::: pysr.BacksolveMutation
+ options:
+ show_root_heading: true
+ heading_level: 3
+ show_root_full_path: false
+
+::: pysr.SimplifyMutation
+ options:
+ show_root_heading: true
+ heading_level: 3
+ show_root_full_path: false
+
+::: pysr.RandomizeMutation
+ options:
+ show_root_heading: true
+ heading_level: 3
+ show_root_full_path: false
+
+::: pysr.OptimizeMutation
+ options:
+ show_root_heading: true
+ heading_level: 3
+ show_root_full_path: false
+
+::: pysr.DoNothingMutation
+ options:
+ show_root_heading: true
+ heading_level: 3
+ show_root_full_path: false
+
+::: pysr.AbstractMutation
+ options:
+ show_root_heading: true
+ heading_level: 3
+ show_root_full_path: false
+
+## Plugins
+
+::: pysr.AdaptiveMutationWeightsPlugin
+ options:
+ show_root_heading: true
+ heading_level: 3
+ show_root_full_path: false
+
+::: pysr.AdaptiveParsimonyPlugin
+ options:
+ show_root_heading: true
+ heading_level: 3
+ show_root_full_path: false
+
+::: pysr.MutationBurstPlugin
+ options:
+ show_root_heading: true
+ heading_level: 3
+ show_root_full_path: false
+
+::: pysr.SimulatedAnnealingPlugin
+ options:
+ show_root_heading: true
+ heading_level: 3
+ show_root_full_path: false
+
+::: pysr.AbstractPlugin
+ options:
+ show_root_heading: true
+ heading_level: 3
+ show_root_full_path: false
diff --git a/docs/src/_index.md b/docs/src/_index.md
index 52587ba3b..df6abef05 100644
--- a/docs/src/_index.md
+++ b/docs/src/_index.md
@@ -16,7 +16,7 @@ hero:
text: View on GitHub
link: https://github.com/astroautomata/PySR
image:
- src: https://ai.damtp.cam.ac.uk/symbolicregression/dev/logo.png
+ src: https://raw.githubusercontent.com/MilesCranmer/PySR_Docs/master/images/pysr-logo.png
alt: PySR
features:
@@ -33,8 +33,8 @@ features:
details: "Customize everything: operators, loss functions, complexity, input types, optimizer, and more."
- icon: 🐍
- title: Python Native
- details: Built for Python with Julia backend for maximum performance and ease of use.
+ title: Familiar Python API
+ details: A scikit-learn interface, with results exported to SymPy, PyTorch, or JAX.
---
If you find PySR useful, please cite the paper [arXiv:2305.01582](https://arxiv.org/abs/2305.01582).
diff --git a/docs/src/backend.md b/docs/src/backend.md
index 801f353cb..3dc4b68cb 100644
--- a/docs/src/backend.md
+++ b/docs/src/backend.md
@@ -33,9 +33,9 @@ The main search code can be found in `src/SymbolicRegression.jl`.
Here are some tips:
-- The documentation for the backend is given [here](https://ai.damtp.cam.ac.uk/symbolicregression/dev/).
+- The documentation for the backend is given [here](https://julia.pysr.ai/dev/).
- Throughout the package, you will often see template functions which typically use a symbol `T` (such as in the string `where {T<:Real}`). Here, `T` is simply the datatype of the input data and stored constants, such as `Float32` or `Float64`. Writing functions in this way lets us write functions generic to types, while still having access to the specific type specified at compilation time.
-- Expressions are stored as binary trees, using the `Node{T}` type, described [here](https://ai.damtp.cam.ac.uk/symbolicregression/dev/types/#SymbolicRegression.CoreModule.EquationModule.Node).
+- Expressions are stored as binary trees, using the `Node{T}` type, described [here](https://julia.pysr.ai/dev/types#SymbolicRegression.CoreModule.EquationModule.Node).
- For reference, the main loop itself is found in the `equation_search` function inside [`src/SymbolicRegression.jl`](https://github.com/astroautomata/SymbolicRegression.jl/blob/master/src/SymbolicRegression.jl).
- Parts of the code which are typically edited by users include:
- [`src/CheckConstraints.jl`](https://github.com/astroautomata/SymbolicRegression.jl/blob/master/src/CheckConstraints.jl), particularly the function `check_constraints`. This function checks whether a given expression satisfies constraints, such as having a complexity lower than `maxsize`, and whether it contains any forbidden nestings of functions.
@@ -70,6 +70,6 @@ For more information on `juliapkg.json`, see [`pyjuliapkg`](https://github.com/J
## Additional notes
-If you get comfortable enough with the backend, you might consider using the Julia package directly: the API is given on the [SymbolicRegression.jl documentation](https://ai.damtp.cam.ac.uk/symbolicregression/dev/).
+If you get comfortable enough with the backend, you might consider using the Julia package directly: the API is given on the [SymbolicRegression.jl documentation](https://julia.pysr.ai/dev/).
If you make a change that you think could be useful to other users, don't hesitate to open a pull request on either the PySR or SymbolicRegression.jl repositories! Contributions are very appreciated.
diff --git a/docs/src/components/VersionPicker.vue b/docs/src/components/VersionPicker.vue
index 17514e206..2dc8a6b49 100644
--- a/docs/src/components/VersionPicker.vue
+++ b/docs/src/components/VersionPicker.vue
@@ -34,6 +34,8 @@ function absoluteUrl(relative) {
const props = defineProps<{ screenMenu?: boolean }>();
const versions = ref>([]);
const currentVersion = ref('Versions');
+// Releases are shown as vX.Y; links still point at the full vX.Y.Z folder.
+const minorLabel = (v: string) => /^(v\d+\.\d+)\.\d+$/.exec(v)?.[1] ?? v;
const isClient = ref(false);
const { site } = useData();
@@ -71,10 +73,21 @@ const loadVersions = async () => {
const scriptsLoaded = await waitForScriptsToLoad();
if (scriptsLoaded && window.DOC_VERSIONS && window.DOCUMENTER_CURRENT_VERSION) {
- versions.value = window.DOC_VERSIONS.map(v => ({
- text: v,
- link: absoluteUrl(`/${v}/`),
- }));
+ // DOC_VERSIONS lists releases newest first, so the first vX.Y.Z of each minor is its latest patch.
+ // Aliases (v1, v1.5) and prereleases stay reachable by URL, out of the picker.
+ const seenMinors = new Set();
+ versions.value = window.DOC_VERSIONS
+ .filter(v => {
+ if (v === 'dev' || v === 'stable') return true;
+ const minor = /^(v\d+\.\d+)\.\d+$/.exec(v)?.[1];
+ if (!minor || seenMinors.has(minor)) return false;
+ seenMinors.add(minor);
+ return true;
+ })
+ .map(v => ({
+ text: minorLabel(v),
+ link: absoluteUrl(`/${v}/`),
+ }));
currentVersion.value = window.DOCUMENTER_CURRENT_VERSION;
} else {
versions.value = [{ text: 'dev', link: absoluteUrl('/dev/') }];
@@ -89,12 +102,12 @@ const loadVersions = async () => {
isClient.value = true;
};
-const versionItems = computed(() => {
- return versions.value.map((v) => ({
- text: v.text,
- link: v.link
- }));
-});
+const versionItems = computed(() => versions.value.map((v) => ({
+ text: v.text,
+ link: v.link,
+ target: '_self',
+ noIcon: true
+})));
onMounted(() => {
if (typeof window !== 'undefined') {
@@ -108,12 +121,12 @@ onMounted(() => {
Data generation code
+
+```python
+import numpy as np
+
+ALPHABET = "ABCDEFGHIJKLMNOPQRSTUVWXYZ"
+PLAIN = "THELAWSOFPHYSICSAREWRITTENINTHELANGUAGEOFMATHEMATICS"
+CIPHER = "".join(ALPHABET[(5 * ALPHABET.index(c) + 8) % 26] for c in PLAIN)
+
+X = np.array([[ALPHABET.index(c)] for c in CIPHER], dtype=object)
+y = np.array([ALPHABET.index(c) for c in PLAIN], dtype=object)
+```
+
+
+
+`Letter` stores a code in `0:25`. Its `string` hook displays constants as letters, such as `V` for code `21`.
+
+```python
+from pysr import PySRRegressor, TypeSpec
+
+spec = TypeSpec(
+ "Letter",
+ fields={"code": "Int"},
+ sample="rng -> Letter(rand(rng, 0:25))",
+ mutate="(rng, value, temperature) -> Letter(mod(value.code + rand(rng, (-1, 1)), 26))",
+ string="value -> string(Char('A' + value.code))",
+)
+```
+
+
+The three operators implement addition, subtraction, and multiplication modulo $26$. Every result remains a `Letter`, and the modulus appears directly in each operator. That built-in alphabet size makes this the easier cipher search.
+
+Use absolute code distance as `elementwise_loss` to reward partial improvements. Zero loss means every ciphertext/plaintext pair matches:
+
+```python
+from pysr import PySRRegressor
+
+model = PySRRegressor(
+ type_spec=spec,
+ operators={
+ 2: [
+ "shift(a::Letter, b::Letter) = Letter(mod(a.code + b.code, 26))",
+ "unshift(a::Letter, b::Letter) = Letter(mod(a.code - b.code, 26))",
+ "mix(a::Letter, b::Letter) = Letter(mod(a.code * b.code, 26))",
+ ]
+ },
+ elementwise_loss="letter_loss(prediction::Letter, target::Letter)::Float64 = abs(prediction.code - target.code)",
+ niterations=20,
+ random_state=0,
+)
+model.fit(X, y)
+print(model.equations_[["complexity", "loss", "equation"]].to_string(index=False))
+```
+
+
+A `loss_function` is useful when scoring needs the dataset as a whole, as in the rollout in [Swinging up a cart-pole](/examples/objectives#swinging-up-a-cart-pole-with-a-rollout-objective) or the period certificate in [Inventing a pseudorandom generator](/examples/objectives#inventing-a-pseudorandom-generator-with-no-target).
+
+One recovered expression is `mix(V, shift(S, x0))`: `V` has code $21$ and `S` has code $18$, giving $21(x_0 + 18) \bmod 26$. Since $21$ is the multiplicative inverse of $5$ modulo $26$, this expression decodes the ciphertext.
+
+### Learning the modulus too
+
+The harder variant removes the alphabet size from the operators. It uses values in `0:51` with integer addition, multiplication, and modulo, so the search must discover the constant $26$ along with the decryption law. This removes a strong structural hint and enlarges the constant search.
+
+One recovered decryption law is:
+
+```
+modulo(mul(add(18, x0), 47), 26)
+```
+
+Run `examples/affine_cipher.py` for the first search or `examples/affine_freemod.py` for the full type, operator, and loss setup of this variant. Allow several minutes for the latter.
+
+## Rediscovering Conway's Game of Life
+
+
+
+Conway's Game of Life can be written as a rule about a cell and its live-neighbour count. This example withholds that rule and supplies only its transition table, asking PySR to reconstruct the logical expression.
+
+The table has $18$ distinct states: two current cell states, each paired with every neighbour count from $0$ through $8$. Repeating the table $8$ times produces $144$ rows for scoring without introducing new cases. A zero-one loss therefore charges $1/18$ for a rule that misses one distinct state. The arrays use `dtype=object` so each value can be wrapped in the custom Julia type.
+
+
+Data generation code
+
+```python
+import numpy as np
+
+TABLE = [(a, n, int(n == 3 or (a == 1 and n == 2))) for a in (0, 1) for n in range(9)]
+REPLICAS = 8
+X = np.array([[a, n] for a, n, _ in TABLE] * REPLICAS, dtype=object)
+y = np.array([t for _, _, t in TABLE] * REPLICAS, dtype=object)
+```
+
+
+
+The `Cell` type wraps a single `Int`, which is enough for the current state, the neighbour count, and the output code. Its hooks propose small counts, move existing constants locally modulo $9$, and print the integer payload.
+
+```python
+from pysr import PySRRegressor, TypeSpec
+
+spec = TypeSpec(
+ "Cell",
+ fields={"v": "Int"},
+ sample="rng -> Cell(rand(rng, 0:8))",
+ mutate="(rng, value, temperature) -> Cell(mod(value.v + rand(rng, (-1, 1)), 9))",
+ string="value -> string(value.v)",
+)
+```
+
+
+The operator vocabulary contains `not`, `and`, `or`, and the equality test `eq`. Arithmetic and threshold comparisons are absent, so a neighbour threshold can enter an expression only as a constant in a call such as `eq(n, 3)`. The loss is zero-one equality on each table row. It records exact agreement, which suits a discrete truth table where the numerical distance between cell codes carries no meaning.
+
+```python
+from pysr import PySRRegressor
+
+model = PySRRegressor(
+ type_spec=spec,
+ operators={
+ 1: ["not(a::Cell) = Cell(a.v == 0 ? 1 : 0)"],
+ 2: [
+ "and(a::Cell, b::Cell) = Cell((a.v != 0 && b.v != 0) ? 1 : 0)",
+ "or(a::Cell, b::Cell) = Cell((a.v != 0 || b.v != 0) ? 1 : 0)",
+ "eq(a::Cell, b::Cell) = Cell(a.v == b.v ? 1 : 0)",
+ ],
+ },
+ elementwise_loss="cell_loss(prediction::Cell, target::Cell)::Float64 = prediction.v == target.v ? 0.0 : 1.0",
+ niterations=160,
+ random_state=0,
+)
+model.fit(X, y, variable_names=["alive", "n"])
+print(model.equations_[["complexity", "loss", "equation"]].to_string(index=False))
+```
+
+
+A recovered rule is:
+
+```
+or(eq(n, 3), and(eq(n, 2), alive))
+```
+
+The first branch allows birth or survival at three neighbours. The second adds survival of a live cell at two neighbours. It agrees with all 18 possible input states, so the transition table fully checks this local rule.
+
+Run `examples/game_of_life.py` to reproduce the Game of Life search.
+
+## Searching over machine words
+
+
+
+This search targets the bit interleave used by Morton, or Z-order, indexing. It gives every expression a `Word` value and restricts the vocabulary to machine instructions: xor, and, or, not, left shift, and right shift. Shift distances are ordinary `Word` constants, so the search must find them as part of the expression.
+
+The reference interleave is evaluated on all $256$ pairs of $4$-bit inputs.
+
+Data generation code
+
+```python
+import numpy as np
+
+from pysr import PySRRegressor, TypeSpec
+
+WIDTH = 4
+WIDER = 5
+
+
+def morton(x, y, width=WIDTH):
+ z = 0
+ for i in range(width):
+ z |= ((x >> i) & 1) << (2 * i)
+ z |= ((y >> i) & 1) << (2 * i + 1)
+ return z
+
+
+PAIRS = [(a, b) for b in range(1 << WIDTH) for a in range(1 << WIDTH)]
+X = np.array(PAIRS, dtype=object)
+y = np.array([morton(a, b) for a, b in PAIRS], dtype=object)
+```
+
+
+
+The `Word` type wraps a `UInt32`. Its `sample` hook favors small values because useful shift distances are small while mask constants still need to reach across the full word. `mutate` either flips a single bit or takes a random unsigned step. This lets a mask change one bit at a time. The `string` hook renders small values as decimals and larger values as zero-padded hexadecimal, keeping shift constants and masks easy to distinguish. The `shl` and `shr` definitions below pass `b.bits` directly; this example does not mask the shift count into `0:31`, so arbitrary sampled `UInt32` counts follow the underlying Julia shift semantics.
+
+```python
+from pysr import TypeSpec
+
+spec = TypeSpec(
+ "Word",
+ fields={"bits": "UInt32"},
+ sample="rng -> Word(rand(rng, Bool) ? rand(rng, UInt32(0):UInt32(31)) : rand(rng, UInt32))",
+ mutate="(rng, value, temperature) -> Word(rand(rng, Bool) ? xor(value.bits, UInt32(1) << rand(rng, 0:31)) : value.bits + rand(rng, (UInt32(1), typemax(UInt32))))",
+ string='value -> value.bits < UInt32(32) ? string(value.bits) : "0x" * string(value.bits, base = 16, pad = 8)',
+)
+```
+
+
+Each operator lifts one machine instruction to `Word`. The loss counts the set bits in the xor of prediction and target and divides by $32$. A zero loss therefore means that every output bit agrees for every scored pair.
+
+```python
+from pysr import PySRRegressor
+
+model = PySRRegressor(
+ type_spec=spec,
+ operators={
+ 1: ["bnot(a::Word) = Word(~a.bits)"],
+ 2: [
+ "bxor(a::Word, b::Word) = Word(xor(a.bits, b.bits))",
+ "band(a::Word, b::Word) = Word(a.bits & b.bits)",
+ "bor(a::Word, b::Word) = Word(a.bits | b.bits)",
+ "shl(a::Word, b::Word) = Word(a.bits << b.bits)",
+ "shr(a::Word, b::Word) = Word(a.bits >> b.bits)",
+ ],
+ },
+ elementwise_loss="bit_loss(prediction::Word, target::Word)::Float64 = count_ones(xor(prediction.bits, target.bits)) / 32",
+ maxsize=45,
+ niterations=800,
+ random_state=0,
+ verbosity=0,
+)
+
+model.fit(X, y, variable_names=["x", "y"])
+```
+
+
+Set `maxsize` large enough for the intended computation. A four-bit reference interleaver uses 41 tree nodes because repeated subexpressions each count toward its size; `maxsize=45` leaves room for that construction.
+
+An expression that matches all four-bit pairs need not work at larger widths: the training domain leaves parts of the masks unconstrained. Use the script's `exact_front_rows` helper to check four-bit agreement and `wider_agreement(model, index)` to test a selected expression on five-bit inputs.
+
+Run `examples/morton_interleave.py` to reproduce this machine-word search.
+
+## Turtle graphics: searching over drawings
+
+
+
+The excerpts below come from `examples/turtle_graphics.py`. Run that script for the complete example; it supplies the Python `path` converter and the Julia `PREAMBLE` containing the underscore-prefixed drawing, mutation, and scoring helpers and constants.
+
+Here the value is a drawing program. A `Path` is a list of pen commands, `forward` and `turn`; operators repeat a path, reflect it, and concatenate command lists. Each expression therefore evaluates to a program, and the value of that program is the picture it draws.
+
+The drawing uses two parallel vectors. `kind=0` marks forward motion. `kind=1` marks a turn. `val` stores the forward length or the turn angle in radians. The `ARM` list below writes turn angles in degrees, and the Python `path` helper converts those entries to radians before storing `y`; the custom `string` hook prints stored turns back in degrees. The target is a plus-sign outline made from four copies of an arm with three edges. The dataset has one row and one feature, `f`, containing a unit forward step. Turns and the path operators must supply the rest of the program.
+
+```python
+import numpy as np
+from pysr import PySRRegressor, TypeSpec
+
+ARM = [("F", 1.0), ("T", 90.0), ("F", 1.0), ("T", -90.0), ("F", 1.0), ("T", 90.0)]
+TARGET = ARM * 4
+
+X = np.empty((1, 1), dtype=object)
+X[0, 0] = ([0], [1.0])
+y = np.empty(1, dtype=object)
+y[0] = path(TARGET)
+VARIABLE_NAMES = ["f"]
+```
+
+
+### Scoring a picture
+
+Two programs can trace the same outline in different orders, so the loss compares rasterized pictures instead of command lists. Each drawing is placed on a $96\times96$ grid after its bounding box is centered and its larger extent is scaled to the grid. The comparison therefore responds to shape while ignoring translation and overall size.
+
+The score is a symmetric chamfer distance between the two sets of lit pixels, measured in pixels and divided by the grid width. The score averages the distance from candidate pixels to the target and from target pixels to the candidate, so zero means the raster masks coincide.
+
+The chamfer loss and the raster
+
+```julia
+# Symmetric chamfer distance in pixels, divided by the grid width. Zero means
+# the two drawings lit exactly the same pixels.
+function _chamfer(pk::Vector{Int8}, pv::Vector{Float64},
+ tk::Vector{Int8}, tv::Vector{Float64})
+ a = _raster(pk, pv)
+ a === nothing && return _FAIL
+ b, db = _target_grids(tk, tv)
+ any(b) || return _FAIL
+ da = _edt(a)
+ ab = _mean_to(a, db); ba = _mean_to(b, da)
+ (ab === nothing || ba === nothing) && return _FAIL
+ v = 0.5 * (ab + ba) / _GRID
+ return isfinite(v) ? v : _FAIL
+end
+```
+
+
+
+### The pen program type
+
+The `TypeSpec` for `Path` stores the two vectors and supplies hooks for continuous turns. `sample` creates one turn at an angle in the interval from $-180$ to $180$ degrees. `mutate` selects a turn and either perturbs it by a Gaussian step with scale $30$ degrees times the temperature or resamples it. The scalar-constant hooks expose every turn to BFGS, which can polish an angle after evolution has placed it near a good value.
+
+Constant folding can collapse an all-constant subtree into a multi-command path. `mutate` and `string` therefore handle values command by command, and the printer represents a folded multi-command constant as `p[...]`.
+
+```python
+from pysr import TypeSpec
+
+PATH = TypeSpec(
+ "Path",
+ fields={"kind": "Vector{Int8}", "val": "Vector{Float64}"},
+ preamble=PREAMBLE,
+ sample="rng -> Path(Int8[1], [_real(rng)])",
+ mutate="""(rng, value, temperature) -> begin
+ turns = findall(isone, value.kind)
+ isempty(turns) && return Path(Int8[1], [_real(rng)])
+ i = rand(rng, turns)
+ val = copy(value.val)
+ val[i] = rand(rng, Bool) ? _perturb(rng, val[i], temperature) : _real(rng)
+ Path(copy(value.kind), val)
+ end""",
+ scalar_constants="value -> _turns(value.kind, value.val)",
+ with_scalar_constants="(value, c) -> Path(copy(value.kind), _put_turns(value.kind, value.val, c))",
+ is_valid="value -> !isempty(value.kind) && all(isfinite, value.val)",
+ string="value -> _literal(value.kind, value.val)",
+)
+```
+
+
+The five operators form a small program algebra. `seq` runs one path after another. `dup` and `tri` repeat a path twice and three times. `mir` negates every turn, reflecting the path. `nest` performs an L-system substitution: it replaces each forward command with a copy of the full path. Each copy is scaled so its net displacement matches the forward command it replaces. Repetition and reflection let a short expression draw many strokes.
+
+```python
+OPERATORS = {
+ 1: [
+ "dup(a::Path) = (t = _seq(a.kind, a.val, a.kind, a.val); Path(t[1], t[2]))",
+ "tri(a::Path) = (t = _seq(a.kind, a.val, a.kind, a.val); u = _seq(t[1], t[2], a.kind, a.val); Path(u[1], u[2]))",
+ "mir(a::Path) = (t = _mir(a.kind, a.val); Path(t[1], t[2]))",
+ "nest(a::Path) = (t = _nest(a.kind, a.val); Path(t[1], t[2]))",
+ ],
+ 2: [
+ "seq(a::Path, b::Path) = (t = _seq(a.kind, a.val, b.kind, b.val); Path(t[1], t[2]))"
+ ],
+}
+
+CHAMFER_LOSS = "chamfer(prediction::Path, target::Path)::Float64 = _chamfer(prediction.kind, prediction.val, target.kind, target.val)"
+```
+
+
+### The search
+
+Pass the path type, operators, and picture loss to the estimator:
+
+```python
+from pysr import PySRRegressor
+
+model = PySRRegressor(
+ type_spec=PATH,
+ operators=OPERATORS,
+ elementwise_loss=CHAMFER_LOSS,
+ niterations=20,
+ verbosity=0,
+ random_state=0,
+)
+
+model.fit(X, y, variable_names=["f"])
+print(model.equations_[["complexity", "loss", "equation"]].to_string(index=False))
+```
+
+
+Scoring each candidate requires rasterization and a distance transform, so even this 20-iteration search can take tens of minutes to hours. Zero loss means the raster masks match; it does not require identical command order or exact turn angles.
+
+One recovered drawing program is:
+
+```
+dup(dup(mir(seq(t-179.6493032061754, tri(seq(p[t-270.04570985370526,t-270.04570985370526,t-270.04570985370526], f))))))
+```
+
+The complete drawing search is in `examples/turtle_graphics.py`.
+
+## Evolving a hopping controller
+
+
+
+Each value in this search is a motor command with three numbers, one for each joint of MuJoCo's hopper. The objective measures performance in simulation. Each candidate expression controls the robot for a fixed-horizon rollout, and its loss is the negated reward from that rollout. The search uses no expert policy, demonstration data, or distilled network; it optimises locomotion directly.
+
+The snippets below describe the recorded search. The runnable script, `examples/hopper_controller.py`, replays its saved controller.
+
+The controller is a `Vec3`, three `Float64` components carried through the tree together:
+
+```python
+from pysr import TypeSpec
+
+TypeSpec(
+ "Vec3",
+ fields={"data": "Vector{Float64}"},
+ sample="rng -> Vec3(randn(rng, 3))",
+ scalar_constants="value -> Float64[value.data[1], value.data[2], value.data[3]]",
+ with_scalar_constants="(value, constants) -> Vec3(collect(constants))",
+ is_valid="value -> length(value.data) == 3 && all(isfinite, value.data)",
+ string='value -> "Vec3([" * join(repr.(Float64.(value.data)), ", ") * "])"',
+ loss_type="Float64",
+)
+```
+
+Each sampled constant contains three independent Gaussian draws, giving each
+motor its own bias. The scalar-constant hooks expose all three components to the
+optimiser so it can tune each motor separately.
+
+The operators are componentwise addition, subtraction and multiplication, and nothing else. No control-specific primitive was supplied:
+
+```python
+{
+ 2: [
+ 'vadd(x, y) = Vec3(x.data .+ y.data)',
+ 'vsub(x, y) = Vec3(x.data .- y.data)',
+ 'vmul(x, y) = Vec3(x.data .* y.data)',
+ ]
+}
+```
+
+### Scoring a controller by simulating it
+
+The eleven features encode the hopper state as vectors. They include the three
+joint angles, angular velocities clipped to $[-10, 10]$ and divided by $10$, and two cyclic permutations
+of each. The torso height, pitch, forward velocity, vertical velocity and
+angular velocity are each broadcast across all three components. The three torso velocities are divided by $5$. A single tree
+computes all three motor commands, so the permuted copies allow each joint to use
+other joints' angles. Before reaching the motors, the final value passes through
+a componentwise `tanh` that limits each command to the actuator range.
+
+Each rollout accumulates the standard Hopper-v5 reward, a survival term plus forward velocity minus a small control cost, and divides by the full horizon, so a fall contributes zero for every step it did not survive. That single number, averaged over eight frozen training starts and negated, is the loss. The Julia loss function is a thin call into the rollout evaluator:
+
+```python
+LOSS_SOURCE = "hopper_rollout_loss(tree, dataset, options) = Main.HopperObjective.loss(tree, dataset, options)"
+```
+
+Because the objective never reads the dataset, `X` and `y` are one row of zero vectors and batching is off. All the information comes from the simulator.
+
+### The recovered controller
+
+Here `A` denotes joint angles, `V` scaled joint angular velocities, and `P` the cyclic permutation `(x0, x1, x2) -> (x1, x2, x0)`. The scalars `p`, `f`, `u`, and `w` denote torso pitch and scaled forward, vertical, and angular velocities:
+
+```
+D = A - 2 P(A) + P(V) - 2 (f + w) + c0
+L = -A - (p + w (P(V) - w)) D - c1
+R = -P(A) - u + w + c2
+action = tanh(P(A) - w + L R)
+```
+
+All products are componentwise, with scalars broadcast. The constant vectors are `c0 = (-0.2477, -4.2268, 8.8397)`, `c1 = (0.0645, 0.6181, 0.3969)`, and `c2 = (0.4974, -1.1921, 1.1808)`, rounded here; the replay script retains their full precision.
+
+### Reproducing the result
+
+Replay the controller on twenty held-out starts with twice the training reset noise.
+
+The replay checks whether the hopper stays upright for eight seconds, travels five metres, and completes at least four contact/liftoff cycles. These criteria check sustained hopping beyond the reward used during training.
+
+The result covers small perturbations of one starting pose in simulation. The evaluation does not test obstacles, hardware, or indefinite hopping.
+
+The saved controller was evaluated with Gymnasium 1.2.2 and MuJoCo 3.3.7. Install `gymnasium[mujoco]==1.2.2` and `mujoco==3.3.7`, then run `python examples/hopper_controller.py`.
diff --git a/docs/src/examples/expression-specifications.md b/docs/src/examples/expression-specifications.md
new file mode 100644
index 000000000..0f240d73a
--- /dev/null
+++ b/docs/src/examples/expression-specifications.md
@@ -0,0 +1,476 @@
+# Expression specifications
+
+## Expression specifications
+
+Use an expression specification when part of the equation is known in advance.
+It fixes the structure around learnable subexpressions while preserving PySR's ordinary
+fitting and prediction interfaces. Inspect template components through `model.equations_`;
+the arbitrary Julia `combine` string does not provide the usual SymPy or LaTeX exports.
+Choose `TemplateExpressionSpec` when the outer form is known and the inner functions remain unknown.
+
+### Template Expressions
+
+Write the outer form in `combine` and name each learnable hole in `expressions`.
+`variable_names` supplies the column names available to `combine`, and each
+hole receives the arguments written in its call. Here `f` receives `x_1` and `x_2`,
+while `g` receives `x_3`, so the search keeps those input roles separate. The target
+form is:
+
+$$ y = \sin(f(x_1, x_2)) + g(x_3) $$
+
+The code below creates the data, declares the template, and fits the model:
+
+```python
+import numpy as np
+
+from pysr import PySRRegressor, TemplateExpressionSpec
+
+# Create data
+X = np.random.randn(1000, 3)
+y = np.sin(X[:, 0] + X[:, 1]) + X[:, 2]**2
+
+# Define template: we want sin(f(x1, x2)) + g(x3)
+template = TemplateExpressionSpec(
+ expressions=["f", "g"],
+ variable_names=["x1", "x2", "x3"],
+ combine="sin(f(x1, x2)) + g(x3)",
+)
+
+model = PySRRegressor(
+ expression_spec=template,
+ binary_operators=["+", "*", "-", "/"],
+ unary_operators=["sin"],
+ maxsize=10,
+)
+model.fit(X, y)
+```
+
+### Parametric Expressions
+
+Add learnable parameter vectors to the template when categories share a formula
+but require different coefficients. This example uses category-specific scale and offset
+values for the target:
+
+$$ y = \alpha \sin(x_1) + \beta $$
+
+Here $\alpha$ is the category-specific scale and $\beta$ is the offset. The
+example has $3$ categories, with $\alpha \in \{1.0, 2.0, 0.5\}$ and
+$\beta \in \{0.1, 1.5, -0.5\}$.
+
+```python
+import numpy as np
+
+from pysr import PySRRegressor, TemplateExpressionSpec
+
+# Create data with 2 features and 3 categories
+X = np.random.uniform(-3, 3, (1000, 2))
+category = np.random.randint(0, 3, 1000)
+
+# Parameters for each category
+offsets = [0.1, 1.5, -0.5]
+scales = [1.0, 2.0, 0.5]
+
+# y = scale[category] * sin(x1) + offset[category]
+y = np.array([
+ scales[c] * np.sin(x1) + offsets[c]
+ for x1, c in zip(X[:, 0], category)
+])
+```
+
+Declare the parameter-vector lengths, optimize their entries during the search,
+and use the selected values in the template. The following block defines that
+parametric expression:
+
+```python
+from pysr import TemplateExpressionSpec
+
+template = TemplateExpressionSpec(
+ expressions=["f"],
+ variable_names=["x1", "x2", "category"],
+ parameters={"p1": 3, "p2": 3}, # One parameter per category
+ combine="f(x1, x2, p1[category], p2[category])"
+)
+```
+
+Append the category as the column listed in `variable_names`; the template indexes
+that column when it selects parameter entries. Python labels start at zero, whereas
+Julia arrays start at one, so shift the labels before stacking them.
+
+```python
+import numpy as np
+
+category_p_one = category + 1
+X_with_category = np.column_stack([X, category_p_one])
+```
+
+Fit the model with the augmented feature matrix, and retain this category column
+for any later `predict` call. The template uses it to choose the corresponding
+parameter entry:
+
+```python
+from pysr import PySRRegressor
+
+model = PySRRegressor(
+ expression_spec=template,
+ binary_operators=["+", "*", "-", "/"],
+ unary_operators=["sin"],
+ maxsize=10,
+)
+model.fit(X_with_category, y)
+
+# Evaluate the fitted model on the same feature/category rows:
+model.predict(X_with_category)
+```
+
+See [Expression Specifications](/api#expression-specifications) for the complete
+reference to this API. The same pattern extends to several template holes and to parameter
+vectors whose entries vary by category.
+
+### Learning multiple outputs jointly
+
+`TemplateExpressionSpec` can couple several scalar holes through one scalar
+residual. Put the observed components in `X`, let the template construct each
+prediction, and return the combined error. This is useful when the components share
+a known outer term.
+
+The target consists of $3$ components. Each contains the shared term $\exp(x_1)$
+and its own second term:
+
+$$\begin{align*}
+y_1 &= \exp(x_1) + x_2^2 \\
+y_2 &= \exp(x_1) + \sin(x_3) \\
+y_3 &= \exp(x_1) + x_1 \cdot x_2
+\end{align*}$$
+
+The data and their noise are generated in the following block:
+
+```python
+import numpy as np
+
+from pysr import PySRRegressor, TemplateExpressionSpec
+
+n = 200
+rstate = np.random.RandomState(0)
+x1 = rstate.uniform(-2, 2, n)
+x2 = rstate.uniform(-2, 2, n)
+x3 = rstate.uniform(-2, 2, n)
+
+# True model with shared component exp(x1):
+y1 = np.exp(x1) + x2**2
+y2 = np.exp(x1) + np.sin(x3)
+y3 = np.exp(x1) + x1 * x2
+
+# Add some noise
+y1 += 0.05 * rstate.randn(n)
+y2 += 0.05 * rstate.randn(n)
+y3 += 0.05 * rstate.randn(n)
+```
+
+`X` contains the three inputs followed by the three observed target components,
+in the same order listed in `variable_names`. The template can therefore use the
+target columns while building its residual:
+
+```python
+import numpy as np
+
+X = np.column_stack([x1, x2, x3, y1, y2, y3])
+```
+
+The template returns the sum of squared errors across the three components for each row.
+
+These sums constrain the predicted components, while the individual `shared` and `f*` terms are not uniquely identifiable. A common function can be added to `shared` and subtracted from every `f*` without changing the residual, so the values in the evaluation block are conditional on recovering the intended decomposition.
+
+```python
+from pysr import TemplateExpressionSpec
+
+spec = TemplateExpressionSpec(
+ expressions=["f1", "f2", "f3", "shared"],
+ variable_names=["x1", "x2", "x3", "y1", "y2", "y3"],
+ combine="""
+ v = shared(x1, x2, x3)
+ y1_predicted = v + f1(x1, x2, x3)
+ y2_predicted = v + f2(x1, x2, x3)
+ y3_predicted = v + f3(x1, x2, x3)
+
+ residuals = (
+ abs2(y1 - y1_predicted) +
+ abs2(y2 - y2_predicted) +
+ abs2(y3 - y3_predicted)
+ )
+
+ residuals
+ """
+)
+```
+
+The fit uses a dummy target and an `elementwise_loss` that returns the template
+value, because `combine` already returns a row-wise squared residual. The estimator
+still receives `dummy_y` to satisfy the fit interface; those zeros do not enter the
+residual:
+
+```python
+import numpy as np
+
+from pysr import PySRRegressor
+
+model = PySRRegressor(
+ expression_spec=spec,
+ binary_operators=["+", "-", "*", "/"],
+ unary_operators=["exp", "sin"],
+ maxsize=20,
+ niterations=50,
+ elementwise_loss="(pred, target) -> pred",
+)
+
+dummy_y = np.zeros(n)
+model.fit(X, dummy_y)
+```
+
+Inspect the fitted named expressions through their Julia objects:
+
+```python
+# Simply get the expression with the highest score:
+idx = model.equations_.score.idxmax()
+
+# Extract the Julia object:
+julia_expr = model.equations_.loc[idx, 'julia_expression']
+
+# Access individual subexpressions:
+for name in ['f1', 'f2', 'f3', 'shared']:
+ tree = getattr(julia_expr.trees, name)
+ print(f"{name}: {tree}")
+```
+
+The named trees can also be evaluated directly through the Julia bridge:
+
+```python
+import numpy as np
+
+from pysr import jl
+from pysr.julia_helpers import jl_array
+
+SR = jl.SymbolicRegression
+
+# Get individual trees
+f1_tree = julia_expr.trees.f1
+shared_tree = julia_expr.trees.shared
+
+# Evaluate one point (x1=1, x2=2, x3=3). `eval_tree_array` expects
+# features by rows, so this is a 3-by-1 array.
+test_inputs = jl_array(np.array([[1.0], [2.0], [3.0]]))
+f1_result, _ = SR.eval_tree_array(f1_tree, test_inputs, model.julia_options_)
+shared_result, _ = SR.eval_tree_array(shared_tree, test_inputs, model.julia_options_)
+
+print(f"f1 at (1,2,3): {f1_result[0]}") # Intended target term: x2^2 = 4.0
+print(f"shared at (1,2,3): {shared_result[0]}") # Intended shared term: exp(1) ≈ 2.718
+```
+
+## Recovering a magnetic field from force measurements
+
+
+
+The force law is fixed in this example, while its coefficient functions are unknown.
+A charged particle in a viscous medium follows
+
+$$ F = -\eta(T)\,v + v \times B(t), $$
+
+The data provide $(t, v, T)$ and the resulting force. The template fixes the drag and
+cross-product assembly and leaves, as learnable holes, the three components of $B(t)$
+and the temperature-dependent drag scale $\eta(T)$.
+
+Each target is a three-component value, so the example defines a `Force` type with three
+fields. Pairing `TypeSpec` with `TemplateExpressionSpec` lets `PySRRegressor` pass these
+custom values through the fit while the template constructs each force from the known law.
+Here the template reads the `.x` field from each returned hole; the other fields remain
+part of the `Force` representation.
+
+The complete runnable script is `examples/magnetic_field.py`; the blocks below expose its
+data and model definitions in the same order.
+
+
+Data generation code
+
+Each scalar input is embedded as a diagonal `Force`, with the same value in each of the
+three fields. The output remains a full three-component force. This lets the template recover
+scalar inputs from `.x` while keeping one custom value type for both `X` and `y`:
+
+`TypeSpec` converts each tuple in these object arrays into the generated `Force` using the
+declared field order (`x`, `y`, `z`), so Python does not need to call `Force(...)` directly.
+
+```python
+import numpy as np
+
+from pysr import PySRRegressor, TemplateExpressionSpec, TypeSpec
+
+OMEGA = 2 * np.pi
+
+
+def experiments(n, seed):
+ rng = np.random.default_rng(seed)
+ T = 298.15 + 0.5 * rng.random(n)
+ t = 10.0 * rng.random(n)
+ v = 2.0 * rng.random((n, 3)) - 1.0
+ B = np.stack([np.sin(OMEGA * t), np.cos(OMEGA * t), np.exp(-t / 10.0)], axis=1)
+ F = -1e-5 * np.sqrt(T)[:, None] * v + np.cross(v, B)
+ return np.column_stack([t, v[:, 0], v[:, 1], v[:, 2], T]), F
+
+
+def as_forces(rows):
+ """An (n, 3) float array as a column of n Force values."""
+ out = np.empty(len(rows), dtype=object)
+ out[:] = [tuple(map(float, row)) for row in rows]
+ return out
+
+
+def on_diagonal(columns):
+ out = np.empty(columns.shape, dtype=object)
+ for j in range(columns.shape[1]):
+ out[:, j] = as_forces(np.repeat(columns[:, j, None], 3, axis=1))
+ return out
+
+
+INPUTS, FORCES = experiments(1000, seed=0)
+X = on_diagonal(INPUTS)
+y = as_forces(FORCES)
+variable_names = ["t", "v_x", "v_y", "v_z", "T"]
+
+HELD_OUT_INPUTS, HELD_OUT_FORCES = experiments(500, seed=12345)
+HELD_OUT_X = on_diagonal(HELD_OUT_INPUTS)
+```
+
+
+
+`FORCE` declares the three fields and the hooks that control sampling, constant
+optimization, string output, and failed evaluation. `sample` draws a value for each
+field. The paired `scalar_constants` and `with_scalar_constants` hooks expose and restore
+all three components, so every `Force` constant is a vector. `init_invalid` supplies a
+`Force` with `NaN` in every field. No explicit `is_valid` is given, so the
+default finite-constant check rejects that value:
+
+```python
+from pysr import TypeSpec
+
+FORCE = TypeSpec(
+ "Force",
+ fields={"x": "Float64", "y": "Float64", "z": "Float64"},
+ sample="rng -> Force(randn(rng, 3)...)",
+ definitions="""
+ Base.sin(a::Force)::Force = Force(sin(a.x), sin(a.y), sin(a.z))
+ Base.cos(a::Force)::Force = Force(cos(a.x), cos(a.y), cos(a.z))
+ function Base.sqrt(a::Force)::Force
+ f(v) = v < 0 ? NaN : sqrt(v)
+ return Force(f(a.x), f(a.y), f(a.z))
+ end
+ Base.exp(a::Force)::Force = Force(exp(a.x), exp(a.y), exp(a.z))
+ Base.:+(a::Force, b::Force)::Force = Force(a.x + b.x, a.y + b.y, a.z + b.z)
+ Base.:-(a::Force, b::Force)::Force = Force(a.x - b.x, a.y - b.y, a.z - b.z)
+ Base.:*(a::Force, b::Force)::Force = Force(a.x * b.x, a.y * b.y, a.z * b.z)
+ Base.:/(a::Force, b::Force)::Force = Force(a.x / b.x, a.y / b.y, a.z / b.z)
+ """,
+ scalar_constants="value -> [value.x, value.y, value.z]",
+ with_scalar_constants="(value, c) -> Force(c[1], c[2], c[3])",
+ string="value -> sprint(show, (value.x, value.y, value.z); context = :compact => true)",
+ init_invalid="() -> Force(NaN, NaN, NaN)",
+)
+```
+
+`definitions` runs after the `Force` type is generated, so its Julia methods can return
+`Force` values. The guarded `sqrt` returns `NaN` for negative fields instead of throwing.
+Register the methods by arity and define a loss over all three force components:
+
+```python
+OPERATORS = {
+ 1: [
+ "Base.sin",
+ "Base.cos",
+ "Base.sqrt",
+ "Base.exp",
+ ],
+ 2: [
+ "Base.:+",
+ "Base.:-",
+ "Base.:*",
+ "Base.:/",
+ ],
+}
+
+FORCE_LOSS = "force_loss(a::Force, b::Force)::Float64 = (a.x - b.x)^2 + (a.y - b.y)^2 + (a.z - b.z)^2"
+```
+
+Inside a template `combine`, PySR's backend wraps input columns and named-hole results in
+`ValidVector` values. The wrapper exposes `.x` for raw values and `.valid` for evaluation
+status; `ValidVector(raw, valid)` constructs a wrapper for a new result. This type is supplied
+by the template runtime rather than defined in this example.
+
+`PHYSICS` checks the hole results for validity, then assembles one `Force` per row
+using the fixed cross product and drag law:
+
+```python
+from pysr import TemplateExpressionSpec
+
+PHYSICS = r"""
+begin
+ _B_x = B_x(t)
+ _B_y = B_y(t)
+ _B_z = B_z(t)
+ _F_d_scale = F_d_scale(T)
+ if !(_B_x.valid && _B_y.valid && _B_z.valid && _F_d_scale.valid)
+ return ValidVector(_B_x.x, false)
+ end
+ F = map(_B_x.x, _B_y.x, _B_z.x, _F_d_scale.x,
+ v_x.x, v_y.x, v_z.x) do bx, by, bz, fd, ux, uy, uz
+ b1, b2, b3 = bx.x, by.x, bz.x
+ u1, u2, u3 = ux.x, uy.x, uz.x
+ s = fd.x
+ Force(u2 * b3 - u3 * b2 + s * u1,
+ u3 * b1 - u1 * b3 + s * u2,
+ u1 * b2 - u2 * b1 + s * u3)
+ end
+ ValidVector(F, true)
+end
+"""
+
+STRUCTURE = TemplateExpressionSpec(
+ combine=PHYSICS,
+ expressions=["B_x", "B_y", "B_z", "F_d_scale"],
+ variable_names=variable_names,
+)
+```
+
+The field holes receive only `t`; `F_d_scale` receives only `T`. Only the `.x` fields
+of their `Force` values enter the numerical residual. The `.y` and `.z` fields can still
+affect whether intermediate values remain valid.
+
+```python
+from pysr import PySRRegressor
+
+model = PySRRegressor(
+ type_spec=FORCE,
+ expression_spec=STRUCTURE,
+ operators=OPERATORS,
+ elementwise_loss=FORCE_LOSS,
+ niterations=100,
+ verbosity=0,
+ random_state=0,
+)
+
+model.fit(X, y, variable_names=variable_names)
+print(model.equations_[["complexity", "loss", "equation"]].to_string(index=False))
+```
+
+An example fitted expression is:
+
+ B_x = sin(#1 / (0.159155, 6.43746, -0.55788))
+ B_y = cos(#1 / (0.159155, 6.43746, -0.55788))
+ B_z = exp(#1 * (-0.1, -0.334652, -1.00446))
+ F_d_scale = -0.000172745, 0.262294, 0.471951
+
+`#1` is the first argument passed to each hole: `t` for the field components and
+`T` for the drag scale. Constants print three entries because they are full `Force`
+values. The field expressions match the sinusoidal and exponential forms used to
+generate the data; the constant drag scale only approximates the temperature-dependent law.
+
+The script validates force predictions on independently generated experiments with
+`held_out_error(model)`. Accurate force predictions alone do not establish exact
+recovery of every coefficient function. Expect an expensive search: a run took about
+1.5 hours in the recorded full-vector configuration.
diff --git a/docs/src/examples/getting-started.md b/docs/src/examples/getting-started.md
new file mode 100644
index 000000000..5b380d4cd
--- /dev/null
+++ b/docs/src/examples/getting-started.md
@@ -0,0 +1,153 @@
+# Getting started
+
+## Simple search
+
+This target combines a cosine term, a square term, and a constant.
+
+The model searches with `+`, `-`, `*`, and `/`. To include the cosine used in the target, add `unary_operators=["cos"]`.
+
+```python
+import numpy as np
+
+from pysr import PySRRegressor
+
+X = 2 * np.random.randn(100, 5)
+y = 2 * np.cos(X[:, 3]) + X[:, 0] ** 2 - 2
+model = PySRRegressor(binary_operators=["+", "-", "*", "/"])
+model.fit(X, y)
+print(model)
+```
+
+## Custom operator
+
+A custom operator has a definition for the symbolic search and a callable for SymPy export and for evaluating equations. Here `inv(x) = 1/x` defines the Julia-side reciprocal operator, while `extra_sympy_mappings` gives that name its Python implementation.
+
+```python
+import numpy as np
+
+from pysr import PySRRegressor
+
+X = 2 * np.random.randn(100, 5)
+y = 1 / X[:, 0]
+model = PySRRegressor(
+ binary_operators=["+", "*"],
+ unary_operators=["inv(x) = 1/x"],
+ extra_sympy_mappings={"inv": lambda x: 1/x},
+)
+model.fit(X, y)
+print(model)
+```
+
+The generic `operators` parameter replaces the separate `unary_operators` and `binary_operators`. Its integer keys specify arity, so a dictionary can describe unary, binary, and higher-arity operators. See the [Operators of any arity](/examples/search-behaviour#operators-of-any-arity) section for this form.
+
+## Multiple outputs
+
+A two-dimensional target makes the model multi-output. Here each column of `y` is the reciprocal of a different input feature, so the fitted model supplies one selected equation for each output.
+
+```python
+import numpy as np
+
+from pysr import PySRRegressor
+
+X = 2 * np.random.randn(100, 5)
+y = 1 / X[:, [0, 1, 2]]
+model = PySRRegressor(
+ binary_operators=["+", "*"],
+ unary_operators=["inv(x) = 1/x"],
+ extra_sympy_mappings={"inv": lambda x: 1/x},
+)
+model.fit(X, y)
+```
+
+## Plotting an expression
+
+After fitting the multi-output model, `model.latex()[0]` returns the LaTeX form of the selected equation for output 0.
+
+```python
+model.latex()[0]
+```
+
+Use `model.latex()[1]` for output 1.
+
+The plot below compares the target and prediction for output 0.
+
+```python
+from matplotlib import pyplot as plt
+plt.scatter(y[:, 0], model.predict(X)[:, 0])
+plt.xlabel('Truth')
+plt.ylabel('Prediction')
+plt.show()
+```
+
+
+
+The `index` argument chooses a candidate equation from `model.equations_`. Pass it to `predict`, `sympy`, or `latex` to inspect or evaluate another candidate. For multiple outputs, provide one equation index per output in output order, then select the desired output column from the predictions.
+
+## Feature selection
+
+Feature selection reduces the candidate input columns before symbolic regression. Set `select_k_features=5` to use a random-forest preprocessor that retains up to five candidates. For structured high-dimensional data, the paper [2006.11287](https://arxiv.org/abs/2006.11287) discusses breaking the problem into smaller pieces before applying PySR.
+
+The target in this example uses features 3 and 19.
+
+```python
+import numpy as np
+
+X = np.random.randn(300, 30)
+y = X[:, 3]**2 - X[:, 19]**2 + 1.5
+```
+
+Configure feature selection in the model:
+
+```python
+from pysr import PySRRegressor
+
+model = PySRRegressor(
+ binary_operators=["+", "-", "*", "/"],
+ unary_operators=["exp"],
+ select_k_features=5,
+)
+```
+
+Fit the model to select the input columns before the search:
+
+```python
+model.fit(X, y)
+```
+
+PySR prints the selected feature names before starting the search, for example:
+
+```text
+Using features ['x3', 'x5', 'x7', 'x19', 'x21']
+```
+
+The names use zero-based column indices. This selection includes both target features, `x3` and `x19`, plus three other candidates. The symbolic search uses only these retained columns.
+
+## Denoising
+
+When per-observation weights are available, pass `weights` to `fit` and define `elementwise_loss` with the matching signature. Without weights, the function accepts `(prediction, target)`. With weights, it accepts `(prediction, target, weight)`. The weighted example uses `elementwise_loss="myloss(x, y, w) = w * (x - y)^2"`. Pass a full custom objective as `loss_function`.
+
+The `denoise=True` option preprocesses the targets with a Gaussian process before symbolic regression. Its predictions become the targets searched by PySR. Because this example passes no `Xresampled`, the predictions are made at the original rows in `X`.
+
+The target combines an exponential term, two linear terms, and random noise:
+
+```python
+import numpy as np
+
+X = np.random.randn(100, 5)
+noise = np.random.randn(100) * 0.1
+y = np.exp(X[:, 0]) + X[:, 1] + X[:, 2] + noise
+```
+
+Create and fit the denoising model:
+
+```python
+from pysr import PySRRegressor
+
+model = PySRRegressor(
+ binary_operators=["+", "-", "*", "/"],
+ unary_operators=["exp"],
+ denoise=True,
+)
+model.fit(X, y)
+print(model)
+```
diff --git a/docs/src/examples/instrumentation.md b/docs/src/examples/instrumentation.md
new file mode 100644
index 000000000..035440646
--- /dev/null
+++ b/docs/src/examples/instrumentation.md
@@ -0,0 +1,227 @@
+# Instrumentation and workflow
+
+## Using TensorBoard for logging
+
+Pass `TensorBoardLoggerSpec` through `logger_spec` to write search progress to TensorBoard. `TensorBoardLoggerSpec` records the model's search hyperparameters and writes scalar summaries for the run. The example uses `min_loss`, which tracks the current best loss, and `pareto_volume`, which summarizes the Pareto front.
+
+```python
+import numpy as np
+
+from pysr import PySRRegressor, TensorBoardLoggerSpec
+
+rstate = np.random.RandomState(42)
+
+# Uniform dist between -3 and 3:
+X = rstate.uniform(-3, 3, (1000, 2))
+y = np.exp(X[:, 0]) + X[:, 1]
+
+# Create a logger that writes to "logs/run*":
+logger_spec = TensorBoardLoggerSpec(
+ log_dir="logs/run",
+ log_interval=10, # Log every 10 iterations
+)
+
+model = PySRRegressor(
+ binary_operators=["+", "*", "-", "/"],
+ logger_spec=logger_spec,
+)
+model.fit(X, y)
+```
+
+Start TensorBoard with the directory that contains the run logs:
+
+```bash
+tensorboard --logdir logs/
+```
+
+## Recording the genealogy of a search
+
+
+
+The complete recorder and post-processing workflow is in `examples/search_trace.py`. `load_trace` and `ancestry` are example-local helpers, not PySR API methods.
+
+Tracing records how a symbolic-regression population changes. Setting `use_tracing=True` activates the optional recorder. PySR loads the required JSON dependency lazily on first use. It then writes JSONL records to `tracing_file`, one for each population at each iteration, with the live members and their mutation, crossover, tuning, and death events. These recorded fields supply the members and parent relationships that the helpers need to reconstruct a graph and one selected path. The helpers work from recorded events and do not expose every historical branch.
+
+With tracing disabled, PySR does not request any trace output or load the tracing dependency.
+
+The target combines one quadratic term, one cosine term, and an offset. The added Gaussian noise has standard deviation equal to $5\%$ of the target's spread:
+
+```python
+import numpy as np
+
+rng = np.random.default_rng(0)
+X = rng.uniform(-3, 3, size=(200, 3))
+y = 2.5382 * np.cos(X[:, 2]) + X[:, 0] ** 2 - 1.5
+y = y + rng.normal(0, 0.05 * y.std(), size=200)
+```
+
+```python
+from pysr import PySRRegressor
+
+model = PySRRegressor(
+ binary_operators=["+", "-", "*", "/"],
+ unary_operators=["cos", "exp"],
+ population_size=27,
+ populations=8,
+ ncycles_per_iteration=90,
+ maxsize=20,
+ niterations=25,
+ random_state=0,
+ use_tracing=True,
+ tracing_file="search_trace.jsonl",
+)
+model.fit(X, y)
+```
+
+Keep tracing runs short: the files grow with the search, and this `load_trace` helper reads the entire file into memory.
+
+
+Trace-reading code
+
+```python
+import json
+import pathlib
+
+def load_trace(trace_path):
+ nodes, edges = {}, []
+ for line in pathlib.Path(trace_path).read_text().splitlines():
+ record = json.loads(line)
+ if record["record_type"] != "iteration":
+ continue
+ for member in record["members"]:
+ nodes.setdefault(member["ref"], {}).update(member)
+ for ref, entry in record["mutations"].items():
+ parent = int(ref)
+ nodes.setdefault(parent, {}).update(
+ {k: entry[k] for k in ("tree", "loss", "cost", "parent") if k in entry}
+ )
+ for event in entry["events"]:
+ kind = event["type"]
+ detail = event.get("mutation") or event.get("details") or {}
+ if kind == "crossover" and event["parent1"] != parent:
+ continue
+ children = (
+ (event["child1"], event["child2"])
+ if kind == "crossover"
+ else (event["child"],) if "child" in event else ()
+ )
+ for child in children:
+ edges.append(
+ dict(
+ parent=parent,
+ child=child,
+ kind=kind,
+ mutation=detail.get("type"),
+ result=detail.get("result", "accept"),
+ mate=event.get("parent2"),
+ )
+ )
+ for edge in edges:
+ if (
+ edge["mate"] is not None
+ and nodes.get(edge["child"], {}).get("parent") == edge["mate"]
+ ):
+ edge["parent"], edge["mate"] = edge["mate"], edge["parent"]
+ return nodes, edges
+```
+
+
+
+`ancestry` selects the lowest-loss recorded member and follows the first edge for each child. It returns one path, which ends when no edge remains or a reference repeats.
+
+```python
+def ancestry(nodes, edges):
+ first_edge = {}
+ for edge in edges:
+ first_edge.setdefault(edge["child"], edge)
+ winner = min(
+ (ref for ref, n in nodes.items() if n.get("loss") is not None),
+ key=lambda ref: nodes[ref]["loss"],
+ )
+ chain, current, seen = [], winner, set()
+ while current in first_edge and current not in seen:
+ seen.add(current)
+ chain.append(first_edge[current])
+ current = first_edge[current]["parent"]
+ chain.reverse()
+ return winner, chain
+```
+
+Trace records label features `x1`, `x2`, and so on by column order, independently of any `variable_names` passed to the search. Thus `x1` refers to `X[:, 0]`, and `x3` refers to `X[:, 2]` in this example.
+
+Use `examples/search_trace.py` for the complete tracing workflow.
+
+## Closing an agent loop with `guesses=`
+
+The `guesses=` keyword accepts a list of expression strings to seed a fit. An external workflow can inspect `model.equations_`, propose an expression, and pass it into a later search. PySR can then optimize its numeric constants.
+
+This example replays a recorded coding-agent proposal from `PROPOSALS`. The custom `agent(front, tried)` helper returns stored expressions; it does not call a live agent or model service. The complete implementation is in `examples/agent_loop_guesses.py`.
+
+The target is relativistic kinetic energy in natural units, with $c=1$ and speeds up to $0.88c$. Gaussian noise has standard deviation $0.15\%$ of the median exact energy. `NOISE_FLOOR` stores its variance for the loop's training-loss stopping rule.
+
+
+Data generation code
+
+```python
+import numpy as np
+
+rng = np.random.default_rng(42)
+mass = rng.uniform(0.5, 3.0, 160)
+velocity = rng.uniform(0.08, 0.88, 160)
+kinetic_energy_exact = mass * (1 / np.sqrt(1 - velocity**2) - 1)
+NOISE_SIGMA = 0.0015 * float(np.median(kinetic_energy_exact))
+kinetic_energy = kinetic_energy_exact + rng.normal(0.0, NOISE_SIGMA, mass.size)
+NOISE_FLOOR = NOISE_SIGMA**2
+
+X = np.stack([mass, velocity], axis=1)
+y = kinetic_energy
+VARIABLE_NAMES = ["mass", "velocity"]
+```
+
+
+
+```python
+from pysr import PySRRegressor
+
+model_kwargs = dict(
+ operators={1: ["sqrt"], 2: ["+", "-", "*", "/"]},
+ niterations=10,
+ verbosity=0,
+)
+```
+
+The proposal uses the known relativistic structure:
+
+```python
+PROPOSALS = ["(mass / sqrt(1 - (velocity * velocity))) - mass"]
+
+def agent(front, tried):
+ print("front handed to the agent:")
+ print(
+ front[["complexity", "loss", "equation"]].to_string(index=False),
+ "\n",
+ flush=True,
+ )
+ return PROPOSALS[len(tried)] if len(tried) < len(PROPOSALS) else None
+```
+
+The loop begins with `guesses=None` and passes each proposal as a one-element list to the next fit. It stops when the minimum training loss reaches `2 * NOISE_FLOOR` or the helper has no proposal left. This threshold is an example-specific stopping rule, not a lower bound on achievable loss.
+
+```python
+from pysr import PySRRegressor
+
+guesses, tried = None, []
+while True:
+ model = PySRRegressor(**model_kwargs, guesses=guesses, random_state=0)
+ model.fit(X, y, variable_names=VARIABLE_NAMES)
+ front = model.equations_
+ if front["loss"].min() <= 2 * NOISE_FLOOR:
+ break
+ proposal = agent(front, tried)
+ if proposal is None:
+ break
+ tried.append(proposal)
+ guesses = [proposal]
+```
+
+Use `examples/agent_loop_guesses.py` for the complete handoff workflow.
diff --git a/docs/src/examples/objectives.md b/docs/src/examples/objectives.md
new file mode 100644
index 000000000..1766b33c8
--- /dev/null
+++ b/docs/src/examples/objectives.md
@@ -0,0 +1,332 @@
+# Objectives and losses
+
+## Custom objectives
+
+`loss_function` gives the objective the candidate's expression tree. Use it when the score needs to manipulate the tree or read dataset-level information. Use `loss_function_expression` when the score needs the complete expression object. Both modes accept Julia source with one of the signatures below.
+
+```julia
+objective(tree_or_expression, dataset, options)
+objective(tree_or_expression, dataset, options, idx=nothing)
+```
+
+With three arguments, the objective sees the dataset given to the current evaluation. Automatic batching makes that dataset the active batch. A method with `idx=nothing` can receive the full dataset and the selected row indices, allowing the objective to slice the rows itself.
+
+The following objective reproduces ordinary unweighted mean squared error, while retaining access to the optional batch indices:
+
+```python
+from pysr import PySRRegressor
+
+objective = """
+function mse_objective(tree, dataset::Dataset{T,L}, options, idx=nothing)::L where {T,L}
+ X = idx === nothing ? dataset.X : dataset.X[:, idx]
+ y = idx === nothing ? dataset.y : dataset.y[idx]
+ prediction, complete = eval_tree_array(tree, X, options)
+ complete || return L(Inf)
+ return sum(abs2, prediction .- y) / length(y)
+end
+"""
+
+model = PySRRegressor(
+ loss_function=objective,
+ binary_operators=["+", "-", "*", "/"],
+)
+```
+
+`eval_tree_array` returns a completion flag with the predictions. Check that flag before using the array and assign an infinite, or suitably large, loss to an incomplete evaluation. Return the dataset's loss type `L`; it must remain real even when the data type `T` is complex.
+
+PySR imports `SymbolicRegression` into Julia before evaluating these source strings, making `Dataset` and `eval_tree_array` available without another import.
+
+A full objective can also attach a different meaning to the stored tree. Here the two children of a binary root supply the numerator and denominator, forming the rational function $P(X)/Q(X)$. The fully qualified `DE.get_child` calls below access the backend's tree interface.
+
+```python
+from pysr import PySRRegressor
+
+objective = """
+function rational_objective(tree, dataset::Dataset{T,L}, options, idx=nothing)::L where {T,L}
+ tree.degree == 2 || return L(Inf)
+ X = idx === nothing ? dataset.X : dataset.X[:, idx]
+ y = idx === nothing ? dataset.y : dataset.y[idx]
+
+ numerator = SymbolicRegression.InterfaceDynamicExpressionsModule.DE.get_child(tree, 1)
+ denominator = SymbolicRegression.InterfaceDynamicExpressionsModule.DE.get_child(tree, 2)
+ p, p_complete = eval_tree_array(numerator, X, options)
+ q, q_complete = eval_tree_array(denominator, X, options)
+ p_complete && q_complete || return L(Inf)
+
+ prediction = p ./ q
+ return sum(abs2, prediction .- y) / length(y)
+end
+"""
+
+model = PySRRegressor(
+ loss_function=objective,
+ binary_operators=["+", "-"],
+)
+```
+
+This objective ignores the root operator and interprets only its two children. The equation table displays the stored tree. `predict` and symbolic export likewise retain the raw tree and cannot automatically reconstruct the quotient. Keep the objective source and apply the child transformation yourself whenever you evaluate the selected equation.
+
+## Writing the objective in Python
+
+The objective itself can live in Python while Julia supplies a thin shim. The search releases Python's global interpreter lock (GIL) while it runs, so the shim reacquires the lock, calls the Python function, and converts its scalar result to the dataset's loss type. The Python function receives the tree, dataset, and options as Julia objects, and it can call back into Julia, including `eval_tree_array`:
+
+```python
+import numpy as np
+
+from pysr import PySRRegressor, jl
+
+
+def python_objective(tree, dataset, options):
+ prediction, completed = jl.SymbolicRegression.eval_tree_array(
+ tree, dataset.X, options
+ )
+ if not completed:
+ return float("inf")
+
+ prediction = np.asarray(prediction)
+ target = np.asarray(dataset.y)
+ return float(np.mean((prediction - target) ** 2))
+
+
+jl.python_objective = python_objective
+
+jl.seval("""
+using PythonCall
+
+function python_objective_shim(
+ tree, dataset::Dataset{T,L}, options
+)::L where {T,L}
+ PythonCall.GIL.@lock begin
+ return pyconvert(L, python_objective(tree, dataset, options))
+ end
+end
+""")
+
+model = PySRRegressor(
+ loss_function="python_objective_shim",
+ precision=64,
+)
+```
+
+`precision=64` makes the data and loss use `Float64`, matching the Python float returned by this objective.
+
+Use this bridge when the objective needs Python-specific logic. Every evaluation crosses the Julia/Python boundary, and concurrent callbacks serialize on the GIL. Keep throughput-sensitive objectives in Julia.
+
+## Swinging up a cart-pole with a rollout objective
+
+
+
+A supervised example gives a target for each row. A control policy has no label of that kind for its force: a force is good only if the subsequent trajectory is good. `loss_function_expression` provides the candidate expression to a rollout objective that can evaluate it throughout the closed loop. This example runs 500 steps from each of 16 training starts and returns the negative mean reward per step. Most of the starts leave the pole hanging down, so the expression must swing it up, catch it, and balance it with one closed-form policy.
+
+### The plant and what the policy sees
+
+The implemented plant uses a $1.0\,\mathrm{kg}$ cart and a $0.1\,\mathrm{kg}$ pole with half-length $0.5\,\mathrm{m}$, gravity $9.8$, and semi-implicit Euler integration with $dt=0.02\,\mathrm{s}$. Its 500-step horizon is therefore $10\,\mathrm{s}$. The policy output is clipped to `[-1, 1]` and multiplied by $10\,\mathrm{N}$. The five inputs are scaled cart position, scaled cart speed, $\sin\theta$, $\cos\theta$, and scaled pole rate. The sine/cosine pair avoids the wrap discontinuity at $\pm\pi$, while the cosine input `c` supplies a direct uprightness signal.
+
+
+Constants and the plant step
+
+```python
+import numpy as np
+
+M_C, M_P, LENGTH, GRAVITY = 1.0, 0.1, 0.5, 9.8
+FORCE_CAP, DT = 10.0, 0.02
+RAIL, V_SCALE, OMEGA_SCALE = 2.4, 4.0, 8.0
+HORIZON = 500 # 10 s of control at dt = 0.02
+
+
+def step(state, force):
+ x, v, theta, omega = state.T
+ total = M_C + M_P
+ sin, cos = np.sin(theta), np.cos(theta)
+ q = (force + M_P * LENGTH * omega**2 * sin) / total
+ theta_dd = (GRAVITY * sin - cos * q) / (LENGTH * (4 / 3 - M_P * cos**2 / total))
+ x_dd = q - M_P * LENGTH * theta_dd * cos / total
+ v, omega = v + DT * x_dd, omega + DT * theta_dd
+ return np.stack([x + DT * v, v, theta + DT * omega, omega], axis=-1)
+```
+
+
+
+Use the initial-state arrays and `observe` helper in `examples/cartpole_objective.py` to construct `X`. Its columns contain the five policy observations, and `y` is a zero placeholder that the objective ignores. The 16 training starts cover hanging, horizontal, and near-upright states; the 64 held-out starts use wider ranges of position, speed, angle, and pole rate.
+
+### The reward
+
+At each step, the objective accumulates the reward
+
+$$ r = 2\cos\theta - 0.05\,x^2 - 0.01\left(\frac{F}{10}\right)^2 - 0.005\left(\frac{F - F_{\text{prev}}}{10}\right)^2 $$
+
+The upright term contributes up to 2 per step. The remaining terms penalize, in order, cart displacement, normalized actuator force, and normalized force changes. Here $F$ is the clipped force and $F_{\mathrm{prev}}$ is the preceding force. The objective returns the negative mean of $r$ over 500 steps and 16 starts, so an upright stationary policy has loss near $-2$.
+
+The complete objective source is defined in `examples/cartpole_objective.py` as the Python string `CARTPOLE_REWARD`. The block below shows its Julia function body and abbreviates the state unpacking, observation-buffer setup, integration step, and divergence check supplied by that script. Treat it as a behavior excerpt, not a standalone definition.
+
+```julia
+function cartpole_reward(ex, dataset::Dataset{T,L}, options)::L where {T,L}
+ n = size(dataset.X, 2)
+ # ... unpack x, v, theta, omega from the observation columns ...
+ for _ in 1:500
+ # ... refill obs from the current state ...
+ raw, ok = eval_tree_array(ex, obs, options)
+ (ok && length(raw) == n && all(isfinite, raw)) || return L(Inf)
+
+ force = 10.0 .* clamp.(raw, -1.0, 1.0)
+ earned += sum(
+ 2.0 .* cos_t .- 0.05 .* x .^ 2 .-
+ 0.01 .* (force ./ 10.0) .^ 2 .-
+ 0.005 .* ((force .- previous) ./ 10.0) .^ 2
+ )
+ previous = force
+ # ... one semi-implicit Euler step, then the divergence check ...
+ end
+ return L(-earned / (500 * n))
+end
+```
+
+Invalid evaluations and divergent trajectories receive a penalty.
+
+### The search
+
+```python
+import sympy
+
+from pysr import PySRRegressor
+
+model = PySRRegressor(
+ operators={
+ 1: ["square", "abs", "tanh"],
+ 2: ["+", "-", "*", "/", "max", "min"],
+ 3: ["ifelse(t, a, b) = t > 0 ? a : b"],
+ },
+ extra_sympy_mappings={
+ "ifelse": lambda t, a, b: sympy.Piecewise((a, t > 0), (b, True))
+ },
+ loss_function_expression=CARTPOLE_REWARD,
+ loss_scale="linear", # rewards make the loss negative, which log scaling forbids
+ maxsize=28,
+ maxdepth=10,
+ parsimony=0.001,
+ niterations=120,
+ populations=32,
+ population_size=64,
+ random_state=0,
+)
+model.fit(X, y, variable_names=["x_n", "v_n", "s", "c", "omega_n"])
+```
+
+`loss_scale="linear"` is required here because logarithmic scaling accepts nonnegative losses, including zero, while this reward can make the loss negative. The operator set supplies arithmetic, `max`, and `min`; the user-defined ternary `ifelse` can switch between swing-up and balancing behavior. `extra_sympy_mappings` translates that custom operator to a SymPy `Piecewise` expression for export. The search uses multithreading, so repeated fits can differ even with the same nominal seed.
+
+### Evaluating a policy
+
+One controller returned by the search is:
+
+```
+(tanh(((0.46341985 * v_n) + (omega_n + (s + omega_n))) / 0.00043365502) - v_n) / 0.97654253
+```
+
+A low training loss alone does not establish successful swing-up from new initial states. In `examples/cartpole_objective.py`, `mean_rewards(model, HELD_OUT_STARTS, index=...)` evaluates a selected equation on the wider held-out set, and `check(model)` requires positive mean reward on every start. Use these rollout checks when choosing a policy.
+
+Allow roughly two hours for the search configuration above. The complete runnable example is `examples/cartpole_objective.py`.
+
+## Inventing a pseudorandom generator with no target
+
+
+
+Ordinary regression supplies `y` and compares every candidate output with a target. This example instead searches for a 32-bit state update $x \mapsto f(x)$ whose iterates have specified properties. Its score comes from how the map behaves, so the objective defines what counts as a good generator without consulting a target array.
+
+The data rows provide 32 random nonzero seeds. `y` is a zero column retained for the estimator interface; the objective never reads it.
+
+```python
+import numpy as np
+
+from pysr import PySRRegressor, TypeSpec
+
+WIDTH = 32
+SEEDS = [int(v) for v in np.random.default_rng(0).integers(1, 1 << WIDTH, size=32)]
+X = np.empty((len(SEEDS), 1), dtype=object)
+for i, seed in enumerate(SEEDS):
+ X[i, 0] = seed
+y = np.zeros(len(SEEDS), dtype=object)
+```
+
+Each state is an unsigned 32-bit word, wrapped in a Julia type. The `shl` and `shr` operators below pass their `UInt32` shift count directly; this example does not mask counts into `0:31`, so arbitrary sampled counts follow Julia's shift semantics and bits shifted out of the 32-bit word are discarded. The operations do not rotate the state. The `sample` hook favors values from 0 through 31 because those are common shift distances, while still sampling arbitrary `UInt32` values. Half of mutations flip a single bit, allowing a local change to a tap constant; the other half adds either 1 or `typemax(UInt32)`. The `string` hook displays values below 32 in decimal and larger words as eight hexadecimal digits.
+
+```python
+from pysr import TypeSpec
+
+SPEC = TypeSpec(
+ "Word",
+ fields={"bits": "UInt32"},
+ sample="rng -> Word(rand(rng, Bool) ? rand(rng, UInt32(0):UInt32(31)) : rand(rng, UInt32))",
+ mutate="(rng, value, temperature) -> Word(rand(rng, Bool) ? xor(value.bits, UInt32(1) << rand(rng, 0:31)) : value.bits + rand(rng, (UInt32(1), typemax(UInt32))))",
+ string='value -> value.bits < UInt32(32) ? string(value.bits) : "0x" * string(value.bits, base = 16, pad = 8)',
+ loss_type="Float64",
+)
+```
+
+The operators form a restricted bitwise instruction set: complement, XOR, AND, OR, left shift, and right shift. The search excludes arithmetic and floating-point operations.
+
+### What the objective scores
+
+`examples/prng_period.py` defines the objective as a Julia source string named `PRNG_LOSS`. It includes the period-certificate implementation used below.
+
+For every candidate, the objective walks 256 steps from each seed and adds six terms that penalize departures from their desired properties:
+
+- `period` is $1 - \log_2(\mathrm{proven})/32$, with `proven` equal to the certified orbit length when certification succeeds.
+- `balance` is the mean squared bias of each bit column away from one half.
+- `autocorr` is the mean squared normalized self-correlation of each bit column at lags 1, 2, 3, and 5.
+- `diffusion` penalizes an average of fewer than four changed output bits when each input bit is flipped on a fixed probe set.
+- `edge` penalizes output bits that respond to fewer than two distinct input-bit flips across that probe set.
+- `shear` measures correlations between bit `b` and shifted bit `b + d` at one- and two-step lags, for nonzero shifts `d` from -3 to 3. It keeps a plain shift from receiving a good score.
+
+### Certifying a period of four billion without walking it
+
+The requested nonzero orbit contains $2^{32}-1 = 4{,}294{,}967{,}295$ states, so a direct walk is impractical. The certificate evaluates the candidate at three groups of inputs: zero, all 32 basis states, and 32 fixed probes. It requires zero to map to zero and the image of each probe to agree with the matrix assembled from the basis images. These finite checks only screen for GF(2) linearity on the supplied evidence; they do not prove global linearity for an arbitrary expression.
+
+A successful certificate gives the candidate's period provided the expression is globally GF(2)-linear. If certification fails, the objective uses the orbit length observed during the finite walk. The period-checking implementation is in `examples/prng_period.py`.
+
+### Settings
+
+The settings below pass the source-defined `PRNG_LOSS` objective to PySR:
+
+```python
+from pysr import PySRRegressor
+
+model = PySRRegressor(
+ type_spec=SPEC,
+ operators={
+ 1: ["bnot(a::Word) = Word(~a.bits)"],
+ 2: [
+ "bxor(a::Word, b::Word) = Word(xor(a.bits, b.bits))",
+ "band(a::Word, b::Word) = Word(a.bits & b.bits)",
+ "bor(a::Word, b::Word) = Word(a.bits | b.bits)",
+ "shl(a::Word, b::Word) = Word(a.bits << b.bits)",
+ "shr(a::Word, b::Word) = Word(a.bits >> b.bits)",
+ ],
+ },
+ loss_function=PRNG_LOSS,
+ batching=False,
+ maxsize=20,
+ populations=8,
+ population_size=30,
+ ncycles_per_iteration=30,
+ niterations=800,
+ random_state=0,
+ verbosity=0,
+)
+model.fit(X, y, variable_names=["x"])
+```
+
+Keep `batching=False` because the objective walks every seed; a minibatch would score a different problem. Each candidate requires a 256-step walk over all 32 seeds. Allow about an hour for this search.
+
+### Results
+
+The script's `check(model)` looks for a front member that passes a separate finite linearity screen and whose inferred GF(2) matrix has order $2^{32}-1$. As with the objective's certificate, this establishes the candidate's period only if the expression is globally GF(2)-linear.
+
+One recovered expression is:
+
+```
+bxor(shl(x, 1), bxor(shr(bxor(shr(x, 1), x), 1), band(x, 0xffffffe3)))
+```
+
+Fixed shifts, XOR, and AND with a constant mask make this expression globally linear by construction. For such a map, matrix order $2^{32}-1$ implies that every nonzero word lies on one orbit of that length. The remaining score terms assess the generator's finite-sample balance, correlations, and diffusion.
+
+The complete runnable example is `examples/prng_period.py`.
diff --git a/docs/src/examples/physics.md b/docs/src/examples/physics.md
new file mode 100644
index 000000000..156c66f27
--- /dev/null
+++ b/docs/src/examples/physics.md
@@ -0,0 +1,183 @@
+# Physics and units
+
+## Dimensional constraints
+
+Dimensional constraints let PySR use physical units during symbolic search. When the units of the features and target are known, pass a unit expression for each feature and one for the target so candidates with mismatched dimensions receive a penalty.
+
+The fit interface follows the unit syntax of `DynamicQuantities.jl`. The following example uses Astropy to sample a mass $M$, a test mass $m$, and a separation $r$, then evaluates Newton's inverse-square law $F = GMm/r^2$.
+
+```python
+import numpy as np
+from astropy import units as u, constants as const
+
+M = (np.random.rand(100) + 0.1) * const.M_sun
+m = 100 * (np.random.rand(100) + 0.1) * u.kg
+r = (np.random.rand(100) + 0.1) * const.R_earth
+G = const.G
+
+F = G * M * m / r**2
+```
+
+The force values span a wide range. This loss compares their magnitudes in logarithmic space and penalizes a wrong sign:
+
+```python
+elementwise_loss = """function loss_fnc(prediction, target)
+ scatter_loss = abs(log((abs(prediction)+1f-20) / (abs(target)+1f-20)))
+ sign_loss = 10 * (sign(prediction) - sign(target))^2
+ return scatter_loss + sign_loss
+end
+"""
+```
+
+`dimensional_constraint_penalty` adds a cost when a candidate fails dimensional analysis:
+
+```python
+from pysr import PySRRegressor
+
+model = PySRRegressor(
+ binary_operators=["+", "-", "*", "/"],
+ unary_operators=["square"],
+ elementwise_loss=elementwise_loss,
+ complexity_of_constants=2,
+ maxsize=25,
+ niterations=100,
+ populations=50,
+ # Amount to penalize dimensional violations:
+ dimensional_constraint_penalty=10**5,
+)
+```
+
+Pass numeric values to `fit`, with matching units in column order through `X_units` and the target unit through `y_units`. Unit strings follow the Julia syntax documented by [DynamicQuantities.jl](https://ai.damtp.cam.ac.uk/dynamicquantities/dev/#Usage).
+
+```python
+import pandas as pd
+
+# Get numerical arrays to fit:
+X = pd.DataFrame(dict(
+ M=M.to("M_sun").value,
+ m=m.to("kg").value,
+ r=r.to("R_earth").value,
+))
+y = F.value
+
+model.fit(
+ X,
+ y,
+ X_units=["Constants.M_sun", "kg", "Constants.R_earth"],
+ y_units="kg * m / s^2"
+)
+```
+
+The loss is nonnegative, and the dimensional penalty is additive. A candidate with total loss below that penalty therefore did not receive the dimensional-violation charge and is dimensionally consistent for the supplied units. The marker `"[⋅]"` denotes a constant whose units remain free. The unit-aware form below illustrates this notation: `M[kg]` supplies mass, so the constant supplies the acceleration units needed for the target:
+
+```julia
+"y[m s⁻² kg] = (M[kg] * 2.6353e-22[⋅])"
+```
+
+Set `dimensionless_constants_only=True` to restrict fitted constants to dimensionless quantities.
+
+## Using differential operators
+
+A [`TemplateExpressionSpec`](/examples/expression-specifications) can include differential operators in its `combine` string. `D` takes the expression to differentiate and an argument index. Here, `D(f, 1)` differentiates the learned function `f` with respect to its first argument, and `df(x)` evaluates that derivative. Matching those values to data turns antiderivative discovery into regression.
+
+The regression target is the integrand $1/(x^2\sqrt{x^2 - 1})$ on $x > 1$. The template searches for a function $f$ whose derivative matches samples of the integrand, leaving an additive constant undetermined.
+
+```python
+import numpy as np
+
+from pysr import PySRRegressor, TemplateExpressionSpec
+
+x = np.random.uniform(1, 10, (1000,)) # Integrand sampling points
+y = 1 / (x**2 * np.sqrt(x**2 - 1)) # Evaluation of the integrand
+
+expression_spec = TemplateExpressionSpec(
+ expressions=["f"],
+ variable_names=["x"],
+ combine="df = D(f, 1); df(x)",
+)
+
+model = PySRRegressor(
+ binary_operators=["+", "-", "*", "/"],
+ unary_operators=["sqrt"],
+ expression_spec=expression_spec,
+ maxsize=20,
+)
+model.fit(x[:, np.newaxis], y)
+```
+
+The target antiderivative is $f(x) = \frac{\sqrt{x^2 - 1}}{x}$, up to an additive constant.
+
+## Discovering a PDE
+
+For a field $u(x,t)$ sampled on a space-time grid, each point becomes one regression row: its field value and spatial derivatives are features, and its time derivative is the target. PySR then searches for the right-hand side in $u_t = f(u, u_x, u_{xx}, \ldots)$.
+
+The following example simulates the viscous Burgers equation $u_t = -u\,u_x + 0.1\,u_{xx}$ on a periodic domain. Measurements of the field can replace the simulation when they are arranged into the same feature and target arrays.
+
+
+Data generation code
+
+```python
+import numpy as np
+
+L, nx, nu = 2 * np.pi, 128, 0.1
+x = np.linspace(0.0, L, nx, endpoint=False)
+dx = L / nx
+k = 2 * np.pi * np.fft.rfftfreq(nx, d=dx)
+
+def d_dx(u, order=1):
+ return np.real(np.fft.irfft((1j * k) ** order * np.fft.rfft(u), n=nx))
+
+def rhs(u):
+ return -u * d_dx(u) + nu * d_dx(u, order=2)
+
+u = -np.sin(x)
+snapshots = [u.copy()]
+dt = 1e-3
+for i in range(1, 4001):
+ k1 = rhs(u)
+ k2 = rhs(u + 0.5 * dt * k1)
+ k3 = rhs(u + 0.5 * dt * k2)
+ k4 = rhs(u + dt * k3)
+ u = u + (dt / 6.0) * (k1 + 2 * k2 + 2 * k3 + k4)
+ if i % 100 == 0:
+ snapshots.append(u.copy())
+
+U = np.array(snapshots) # shape (41, 128)
+t = np.arange(U.shape[0]) * 0.1
+```
+
+
+
+Estimate spatial derivatives with the smoothing differentiator `savgol_filter` and time derivatives with `np.gradient`. Supply the coordinate spacings and keep each space-time point aligned across the feature columns and target:
+
+```python
+import numpy as np
+
+from scipy.signal import savgol_filter
+
+ux = savgol_filter(U, 21, 3, deriv=1, delta=dx, axis=-1)
+uxx = savgol_filter(U, 21, 3, deriv=2, delta=dx, axis=-1)
+ut = np.gradient(U, t, axis=0)
+
+X = np.stack([U, ux, uxx], axis=-1).reshape(-1, 3)
+y = ut.reshape(-1)
+```
+
+`complexity_of_variables=[1, 2, 3]` makes higher spatial derivatives more expensive, biasing the search toward lower-order terms when candidate fits compete. A successful low-complexity search should place $0.1\,u_{xx} - u\,u_x$ on the Pareto front, matching the right-hand side that generated the snapshots.
+
+```python
+from pysr import PySRRegressor
+
+model = PySRRegressor(
+ binary_operators=["+", "-", "*"],
+ complexity_of_variables=[1, 2, 3],
+ maxsize=20,
+ niterations=100,
+)
+model.fit(X, y, variable_names=["u", "u_x", "u_xx"])
+print(model)
+```
+
+The regression uses estimated derivatives, so simulate the discovered PDE forward from a held-out initial condition before relying on its coefficients.
+
+If the field $u$ has units $\mathrm{m\,s^{-1}}$, then $u_x$ has units $\mathrm{s^{-1}}$, $u_{xx}$ has units $\mathrm{m^{-1}\,s^{-1}}$, and $u_t$ has units $\mathrm{m\,s^{-2}}$. For this feature order, pass the corresponding strings through `X_units=["m/s", "s^-1", "m^-1*s^-1"]` and `y_units="m*s^-2"`.
diff --git a/docs/src/examples/search-behaviour.md b/docs/src/examples/search-behaviour.md
new file mode 100644
index 000000000..e46f0cb7f
--- /dev/null
+++ b/docs/src/examples/search-behaviour.md
@@ -0,0 +1,270 @@
+# Search behaviour
+
+## Automatic batching on a large dataset
+
+`batching` reduces the rows used to score population members during evolution. Hall-of-fame comparisons continue to use every row. With `batching=False`, every evolutionary comparison uses all rows. When `batch_size` is `None`, automatic batching can use every row on small datasets. PySR chooses all rows through 1000, 128 below 5000, 256 below 50,000, and 512 from 50,000 onward. At 20,000 rows, evolution therefore sees 256 rows per comparison, while hall-of-fame scoring sees all 20,000. `batching="auto"` is the default; this example names it to make the choice explicit.
+
+The target depends on two of five input columns. The remaining three inputs act as distractors.
+
+
+Data generation code
+
+```python
+import numpy as np
+
+rng = np.random.default_rng(0)
+X = rng.uniform(-3, 3, (20000, 5))
+y = 2.5382 * np.cos(X[:, 3]) + X[:, 0] ** 2 - 0.5
+```
+
+
+
+```python
+from pysr import PySRRegressor
+
+model = PySRRegressor(
+ binary_operators=["+", "-", "*", "/"],
+ unary_operators=["cos", "exp"],
+ batching="auto",
+ niterations=40,
+ verbosity=0,
+ random_state=0,
+)
+model.fit(X, y, variable_names=["x0", "x1", "x2", "x3", "x4"])
+print(model.equations_[["complexity", "loss", "equation"]].to_string(index=False))
+```
+
+One recovered expression is:
+
+```
+(cos(x3) * 2.5382) + ((x0 * x0) + -0.5)
+```
+
+Run the complete example with `examples/automatic_batching.py`.
+
+## Operators of any arity
+
+
+
+The generic `operators` argument defines an expression grammar by arity. Each integer key gives the number of arguments accepted by the operators listed under it, so one search can mix unary and binary operations with ternary or wider functions. `operators` cannot be combined with `binary_operators` or `unary_operators`.
+
+The target is a four-way minimum,
+
+$$
+y = \min(x + 2,\; 2 - x,\; \sin(3x) + 1.2,\; 0.8).
+$$
+
+A binary `min` would write this reduction as three nested calls. The custom `min4` writes the same mathematical operation in one call, giving the search a shorter structural route to the target.
+
+
+Data generation code
+
+```python
+import numpy as np
+
+DOMAIN = (-4.0, 4.0)
+
+
+def envelope(x):
+ return np.minimum.reduce(
+ [x + 2.0, 2.0 - x, np.sin(3.0 * x) + 1.2, np.full_like(x, 0.8)]
+ )
+
+
+_x = np.sort(np.random.default_rng(0).uniform(*DOMAIN, 400))
+X = _x.reshape(-1, 1).astype(np.float32)
+y = envelope(_x).astype(np.float32)
+
+_unseen = np.linspace(*DOMAIN, 4001)
+X_UNSEEN = _unseen.reshape(-1, 1).astype(np.float32)
+Y_UNSEEN = envelope(_unseen)
+```
+
+
+
+Custom operator definitions are Julia source strings under their arity keys. A SymPy-capable model also needs a matching entry in `extra_sympy_mappings`; PySR rejects a named custom function during setup when its symbolic image is missing. The mappings here provide the `Piecewise` representation for `ifelse` and the `Min` representation for `min4`.
+
+```python
+import sympy
+
+from pysr import PySRRegressor
+
+model = PySRRegressor(
+ operators={
+ 1: ["sin", "cos"],
+ 2: ["+", "-", "*"],
+ 3: ["ifelse(t, a, b) = t > 0 ? a : b"],
+ 4: ["min4(a, b, c, d) = min(a, b, c, d)"],
+ },
+ # A custom operator is rejected at fit time unless it has a SymPy image.
+ extra_sympy_mappings={
+ "ifelse": lambda t, a, b: sympy.Piecewise((a, t > 0), (b, True)),
+ "min4": lambda a, b, c, d: sympy.Min(a, b, c, d),
+ },
+ niterations=500,
+ random_state=0,
+)
+model.fit(X, y, variable_names=["x"])
+print(model.equations_)
+```
+
+One recovered expression is:
+
+```
+min4(0.8, 2.0 - x, 2.0 + x, 1.1999999 + sin(x * 3.0))
+```
+
+The complete example, `examples/any_arity.py`, checks predictions on an unseen grid with `model.predict(X_UNSEEN, index=i)` for each candidate equation.
+
+## Mutations and plugins
+
+PySR accepts mutation and plugin configuration as `PySRRegressor` arguments. `mutations` maps mutation entries to weights, and `plugins` lists plugin entries. Entries merge with shipped defaults by concrete type: a matching type replaces the default entry, while a new type is appended alongside it. This lets one behavior change without reproducing the complete default set.
+
+The following configuration activates Backsolve and names the adaptive mutation-weight plugin.
+
+```python
+from pysr import (
+ AdaptiveMutationWeightsPlugin,
+ BacksolveMutation,
+ PySRRegressor,
+)
+
+model = PySRRegressor(
+ mutations={BacksolveMutation(): 0.1},
+ plugins=[AdaptiveMutationWeightsPlugin()],
+)
+```
+
+`BacksolveMutation` is present at default weight 0.0, so it is available without being sampled. Assigning weight 0.1 makes it eligible for this search. Backsolve inverts the operators above a selected subtree and fits a replacement from expressions already in the population, helping an additive target assemble from separately discovered pieces.
+
+`AdaptiveMutationWeightsPlugin` is part of the default plugin set. It updates mutation weights from observed improvement rates and resets its runtime state for each `fit`. Naming it in `plugins` replaces the existing entry of that type, which is also where `smoothing`, `floor`, and `reward` can be changed. The resolved list keeps simulated annealing and adaptive parsimony alongside adaptive mutation weights. The two settings work together: the mutation weight makes Backsolve eligible, and the adaptive plugin can change how much of the mutation budget each kind receives.
+
+The runnable configuration uses the small target
+
+$$ y = 2.5\cos(3x) + 0.5x^2 - 1. $$
+
+```python
+import numpy as np
+
+x = np.linspace(-3.0, 3.0, 200)
+X = x.reshape(-1, 1)
+y = 2.5 * np.cos(3.0 * x) + 0.5 * x * x - 1.0
+```
+
+```python
+from pysr import AdaptiveMutationWeightsPlugin, BacksolveMutation, PySRRegressor
+
+model = PySRRegressor(
+ binary_operators=["+", "-", "*"],
+ unary_operators=["cos", "exp"],
+ niterations=20,
+ verbosity=0,
+ mutations={BacksolveMutation(): 0.1},
+ plugins=[AdaptiveMutationWeightsPlugin()],
+ random_state=0,
+)
+model.fit(X, y, variable_names=["x"])
+print(model.equations_)
+```
+
+After fitting, inspect `model.julia_options_.mutations` to see the resolved entries and weights. Run the example with `examples/mutations_and_plugins.py`.
+
+## Adaptive mutation weights
+
+
+
+`AdaptiveMutationWeightsPlugin` adjusts mutation weights from their observed success rates. It is enabled by default, and its learned state resets on each `fit`. Pass an instance through `plugins` to change its settings.
+
+- `smoothing` controls the exponential moving average of success rates.
+- `floor` limits the success-rate ratio used to update a multiplier to $[\text{floor}, 1/\text{floor}]$.
+- `reward` selects whether improvement is measured by `"cost"` or `"loss"`.
+
+This example fits a cosine and a quadratic term:
+
+
+Data generation code
+
+```python
+import numpy as np
+
+rng = np.random.default_rng(20260817)
+X = rng.uniform(-3.0, 3.0, size=(200, 5))
+y = 2.5382 * np.cos(X[:, 3]) + X[:, 0] ** 2 - 0.5
+```
+
+
+
+```python
+from pysr import AdaptiveMutationWeightsPlugin, PySRRegressor
+
+model = PySRRegressor(
+ binary_operators=["+", "-", "*", "/"],
+ unary_operators=["cos", "exp"],
+ maxsize=14,
+ niterations=300,
+ populations=1,
+ plugins=[AdaptiveMutationWeightsPlugin(smoothing=0.02, floor=0.05, reward="cost")],
+ random_state=0,
+ verbosity=0,
+)
+model.fit(X, y)
+print(model.equations_)
+```
+
+For an optional diagnostic that records mutation counts and learned multipliers, see `examples/adaptive_mutation_weights.py`.
+
+## The backsolve mutation
+
+
+
+Ordinary mutations change an expression and test the result. Backsolve chooses a non-root subtree and inverts the operators above it to determine the values that subtree would need to produce. It then fits a sparse replacement from expressions already in the population. Greedy forward selection over a basis library finds the replacement within the remaining complexity budget. The replaced subtree can remain in the library, and rows with invalid inverse values are masked. Both `+` and `*` are required for the weighted sum. An event returns a child only when its cost strictly improves on the parent. Each event requires an inversion and a linear solve, so enabling Backsolve can substantially increase search time.
+
+Backsolve applies to ordinary binary scalar expression trees with a real target vector. It rejects unsupported wrappers, shared nodes, root or single-node choices, missing targets, and missing weighted-sum operators. The mutation is experimental and starts at default weight 0.0.
+
+The target combines an oscillation with a quadratic:
+
+$$ y = 2.5\cos(3x) + 0.5x^2 - 1. $$
+
+```python
+import numpy as np
+
+x = np.linspace(-3.0, 3.0, 200)
+X = x.reshape(-1, 1)
+y = 2.5 * np.cos(3.0 * x) + 0.5 * x * x - 1.0
+```
+
+The example omits division and a power operator, so the quadratic must appear as `x * x`. `precision=64` supports the linear solve, and `maxsize=30` leaves room for the complete form.
+
+```python
+from pysr import PySRRegressor
+
+model = PySRRegressor(
+ binary_operators=["+", "-", "*"],
+ unary_operators=["cos", "exp"],
+ weight_backsolve=1.0,
+ maxsize=30,
+ precision=64,
+ niterations=10,
+ verbosity=0,
+ random_state=0,
+)
+model.fit(X, y, variable_names=["x"])
+print(model.equations_[["complexity", "loss", "equation"]].to_string(index=False))
+```
+
+One recovered expression is:
+
+```
+((cos((x + x) + x) * 2.5) + ((x * x) * 0.5)) + -1.0
+```
+
+Check the selected equation against the training target:
+
+```python
+import numpy as np
+
+prediction = np.asarray(model.predict(X), dtype=float)
+print(f"max abs error: {np.max(np.abs(prediction - y)):.3g}")
+```
+
+Use `examples/backsolve_mutation.py` for the complete example.
diff --git a/docs/src/examples/value-types.md b/docs/src/examples/value-types.md
new file mode 100644
index 000000000..53b679410
--- /dev/null
+++ b/docs/src/examples/value-types.md
@@ -0,0 +1,413 @@
+# Value types
+
+## Complex numbers
+
+To fit complex-valued expressions, pass feature and target arrays with a complex NumPy dtype.
+
+The code constructs a complex cosine target and restricts the search to `+`, `-`, `*`, and `cos`:
+
+```python
+import numpy as np
+
+from pysr import PySRRegressor
+
+X = np.random.randn(100, 1) + 1j * np.random.randn(100, 1)
+y = (1 + 2j) * np.cos(X[:, 0] * (0.5 - 0.2j))
+
+model = PySRRegressor(
+ binary_operators=["+", "-", "*"], unary_operators=["cos"], niterations=100,
+)
+
+model.fit(X, y)
+```
+
+The fitted constants belong to the complex-valued expression domain. Use `model.sympy()` to inspect the selected equation in SymPy form:
+
+```python
+model.sympy()
+```
+
+Pass an equation-row index to `predict`; `-1` selects the final row of `equations_`:
+
+```python
+model.predict(X, -1)
+```
+
+## Julia packages and types
+
+PySR delegates its symbolic-regression search to the pure Julia package [SymbolicRegression.jl](https://github.com/astroautomata/SymbolicRegression.jl). A custom operator can call functionality from another Julia package through that backend.
+
+Load the package that supplies the prime lookup before defining the operator. The target in this example is
+
+$$ y = p_{3x + 1} - 5, $$
+
+where $p_i$ denotes the $i$th prime and $x$ is the input feature. The lookup comes from [Primes.jl](https://github.com/JuliaMath/Primes.jl).
+
+First, import the handle to PySR's Julia runtime:
+
+```python
+from pysr import jl
+```
+
+`jl` is the juliacall runtime object, and `jl.seval` evaluates Julia source from Python. Use it to add `Primes.jl` to the active PySR environment:
+
+```python
+from pysr import jl
+
+jl.seval("""
+import Pkg
+Pkg.add("Primes")
+""")
+```
+
+Adding a package and importing its namespace are separate steps. Import `Primes` before defining `p`:
+
+```python
+from pysr import jl
+
+jl.seval("import Primes")
+```
+
+With the package available, define the unary operator used by the search:
+
+```python
+from pysr import jl
+
+jl.seval("""
+function p(i::T) where T
+ if (0.5 < i < 1000)
+ return T(Primes.prime(round(Int, i)))
+ else
+ return T(NaN)
+ end
+end
+""")
+```
+
+The operator rounds inputs within `0.5 < i < 1000` to prime indices and preserves the input type `T`. Outside that domain it returns `T(NaN)` to reject invalid evaluations.
+
+Call the Julia function through `jl.p` to build the prime values used by the synthetic data:
+
+```python
+from pysr import jl
+
+primes = {i: jl.p(i*1.0) for i in range(1, 999)}
+```
+
+Use the lookup table to generate noisy samples of the target:
+
+```python
+import numpy as np
+
+X = np.random.randint(0, 100, 100)[:, None]
+y = [primes[3*X[i, 0] + 1] - 5 + np.random.randn()*0.001 for i in range(100)]
+```
+
+Register `p` with the regular PySR operator interface. The `sympy_p` subclass and `extra_sympy_mappings` give SymPy a symbolic form for the custom function:
+
+```python
+from pysr import PySRRegressor
+import sympy
+
+class sympy_p(sympy.Function):
+ pass
+
+model = PySRRegressor(
+ binary_operators=["+", "-", "*", "/"],
+ unary_operators=["p"],
+ niterations=100,
+ extra_sympy_mappings={"p": sympy_p}
+)
+```
+
+Fit the model on the generated pairs:
+
+```python
+model.fit(X, y)
+```
+
+Since `p` rounds its argument, different real offsets can produce the same integer prime index, so the offset in a fitted expression need not equal `1` exactly. Inspect the selected equation in SymPy with:
+
+```python
+model.sympy()
+```
+
+## Custom value types
+
+`TypeSpec` declares a custom Julia value that can flow through an expression tree. Vectors, tensors, strings, structs, and other Julia value types can serve as payloads when their behavior is described by the type's hooks, operators, and loss.
+
+A model configured with `type_spec=TypeSpec(...)` needs three pieces:
+
+1. Define hooks that sample and change values. For continuously optimized constants, `scalar_constants` exposes scalar coordinates and `with_scalar_constants` rebuilds the value. A custom `mutate` hook handles changes that need a discrete or structural operation.
+2. Choose exactly one of `elementwise_loss`, `loss_function`, or `loss_function_expression`. The selected loss must accept the custom values and return a real-valued score. Elementwise-loss return types are inferred; a full objective requires an explicit `loss_type`.
+3. Supply operators through `operators={...}`, grouped by arity. Each operator must accept and return the declared value type so the type-stability check succeeds. Add an explicit return annotation such as `::Vec2` when Julia cannot infer it.
+
+The examples below combine these pieces for vector values, strings as discrete constants, and payloads with several shapes.
+
+### Vector-valued expression trees
+
+This example searches for an expression whose values are two-dimensional vectors:
+
+$$
+y = \operatorname{rotate90}(x_1) + 2x_2 +
+\begin{bmatrix}0.5 \\ -1.0\end{bmatrix}.
+$$
+
+The 128-by-2 DataFrame has two vector-valued features per row, and each entry of `y` is one target vector. The table dimensions describe samples and features, while each cell of the object array holds a single logical vector value.
+
+```python
+import numpy as np
+import pandas as pd
+
+from pysr import PySRRegressor, TypeSpec
+
+rng = np.random.default_rng(0)
+x1 = [rng.normal(size=2) for _ in range(128)]
+x2 = [rng.normal(size=2) for _ in range(128)]
+X = pd.DataFrame({"x1": x1, "x2": x2})
+y = np.empty(128, dtype=object)
+offset = np.array([0.5, -1.0])
+y[:] = [np.array([-a[1], a[0]]) + 2 * b + offset for a, b in zip(x1, x2)]
+```
+
+Declare `Vec2` with one `Vector{Float64}` field. `sample` creates a random payload of length two, while `scalar_constants` and `with_scalar_constants` expose and rebuild the entries used by continuous constant optimization:
+
+```python
+from pysr import TypeSpec
+
+type_spec = TypeSpec(
+ "Vec2",
+ fields={"data": "Vector{Float64}"},
+ sample="rng -> Vec2(randn(rng, 2))",
+ scalar_constants="value -> value.data",
+ with_scalar_constants=(
+ "(value, scalar_constants) -> Vec2(scalar_constants)"
+ ),
+)
+```
+
+PySR derives how values are initialized, mutated, validated, counted, packed, and unpacked from these hooks. Add `init` or `mutate` for specialized search behavior, or `string` for custom display, when the derived defaults are insufficient.
+
+`TypeSpec` generates the Julia value type from that declaration. Its equivalent Julia struct has one vector field:
+
+```julia
+struct Vec2
+ data::Vector{Float64}
+end
+```
+
+Define operators that consume and return `Vec2` values:
+
+```python
+operators = {
+ 1: [
+ "rotate90(a::Vec2) = Vec2([-a.data[2], a.data[1]])",
+ "double(a::Vec2) = Vec2(2a.data)",
+ ],
+ 2: ["add_vectors(a::Vec2, b::Vec2) = Vec2(a.data + b.data)"],
+}
+```
+
+The elementwise loss compares vector payloads by summing squared differences, which gives the real-valued score the search requires:
+
+```python
+from pysr import PySRRegressor
+
+model = PySRRegressor(
+ type_spec=type_spec,
+ operators=operators,
+ elementwise_loss="vector_loss(a::Vec2, b::Vec2) = sum(abs2, a.data - b.data)",
+ niterations=40,
+ populations=4,
+ maxsize=10,
+)
+
+model.fit(X, y)
+print(model.equations_)
+```
+
+The search space can represent the target as `add_vectors(add_vectors(rotate90(x1), double(x2)), Vec2([0.5, -1.0]))`. The vector constant carries both offset components, which the scalar hook pair exposes for joint optimization.
+
+
+
+String-valued expressions and discrete constants
+
+This example treats each string as a value and chooses a separator from a finite set during mutation. It lowercases the first input, uppercases the second, and joins them with the selected separator:
+
+```python
+import numpy as np
+import pandas as pd
+
+from pysr import PySRRegressor, TypeSpec
+
+X = pd.DataFrame(
+ {
+ "first": ["Py", "symbolic", "hello", "left"],
+ "second": ["SR", "regression", "world", "right"],
+ }
+)
+y = np.array(
+ [f"{a.lower()}-{b.upper()}" for a, b in X.itertuples(index=False)],
+ dtype=object,
+)
+
+type_spec = TypeSpec(
+ "StringValue",
+ fields={"data": "String"},
+ sample='rng -> StringValue(rand(rng, ("", "-", "_")))',
+ mutate="""
+ mutate_string(rng, value, temperature) = StringValue(rand(rng, ("", "-", "_")))
+ """,
+)
+
+model = PySRRegressor(
+ type_spec=type_spec,
+ operators={
+ 1: [
+ "string_lowercase(x::StringValue) = StringValue(lowercase(x.data))",
+ "string_uppercase(x::StringValue) = StringValue(uppercase(x.data))",
+ ],
+ 2: [
+ "string_concat(a::StringValue, b::StringValue) = StringValue(a.data * b.data)"
+ ],
+ },
+ elementwise_loss="""
+ string_loss(a::StringValue, b::StringValue) = Float64(Base.editdistance(a.data, b.data))
+ """,
+ niterations=40,
+)
+
+model.fit(X, y)
+print(model.equations_)
+```
+
+The `sample` and `mutate` hooks choose discrete separator constants. Without a
+`scalar_constants` and `with_scalar_constants` pair, BFGS leaves them unchanged.
+Edit distance provides the real-valued loss for string outputs.
+
+
+
+
+Advanced: recovering a neural network with tensor constants
+
+A `TypeSpec` can place scalar, vector, and matrix payloads in the same Julia value type. Its scalar-constant hooks flatten each payload for BFGS and rebuild its original shape.
+
+The advanced example searches for this two-layer neural network:
+
+$$ y = W_2\operatorname{relu}(W_1x + b_1) + b_2 $$
+
+The targets are vectors. `safe_matmul` and `safe_add` return `NaN` for incompatible shapes, so expressions using those helpers receive an invalid value instead of a dimension error:
+
+```python
+import numpy as np
+import pandas as pd
+
+from pysr import PySRRegressor, TypeSpec
+
+preamble = """
+const NNPayload = Union{Float64, Vector{Float64}, Matrix{Float64}}
+
+safe_matmul(a::Matrix{Float64}, b::Vector{Float64}) =
+ size(a, 2) == length(b) ? a * b : NaN
+safe_matmul(::NNPayload, ::NNPayload) = NaN
+
+safe_add(a::Float64, b::Float64) = a + b
+safe_add(a::T, b::T) where {T<:Union{Vector{Float64}, Matrix{Float64}}} =
+ size(a) == size(b) ? a + b : NaN
+safe_add(::NNPayload, ::NNPayload) = NaN
+
+# Constants sample a random rank, generating only the payload that was chosen:
+function random_nn_payload(rng)
+ rank = rand(rng, 0:2)
+ rank == 0 ? randn(rng) : rank == 1 ? randn(rng, 2) : randn(rng, 2, 2)
+end
+"""
+
+type_spec = TypeSpec(
+ "NNValue",
+ fields={"data": "NNPayload"},
+ sample="rng -> NNValue(random_nn_payload(rng))",
+ # Mutations usually perturb every scalar in the payload, but occasionally
+ # resample a fresh rank:
+ mutate="""
+ (rng, value, temperature) -> if rand(rng) < 0.1
+ NNValue(random_nn_payload(rng))
+ else
+ NNValue(value.data .+ temperature .* randn(rng, size(value.data)...))
+ end
+ """,
+ scalar_constants="""
+ function scalar_constants(value)
+ return value.data isa Float64 ? [value.data] : vec(value.data)
+ end
+ """,
+ with_scalar_constants="""
+ function with_scalar_constants(value, scalar_constants)
+ data = value.data isa Float64 ? scalar_constants[1] :
+ reshape(collect(scalar_constants), size(value.data))
+ return NNValue(data)
+ end
+ """,
+ preamble=preamble,
+)
+```
+
+`preamble` makes the payload union and helpers available before the generated type is
+defined. Mutation can change a constant's rank; continuous optimization preserves its
+shape through the scalar hook pair.
+
+Generate vector-valued targets from fixed weights and biases:
+
+```python
+import numpy as np
+import pandas as pd
+
+rng = np.random.default_rng(0)
+x_values = rng.normal(size=(64, 2))
+W1 = np.array([[1.2, -0.7], [0.5, 1.1]])
+b1 = np.array([0.3, -0.2])
+W2 = np.array([[0.8, -1.0], [1.3, 0.4]])
+b2 = np.array([-0.4, 0.2])
+y_values = (W2 @ np.maximum(x_values @ W1.T + b1, 0).T).T + b2
+
+X = pd.DataFrame({"x": list(x_values)})
+y = pd.Series(list(y_values), dtype=object)
+```
+
+Register the neural-network operations and score matching vector outputs with mean squared error. The loss penalizes other output shapes:
+
+```python
+from pysr import PySRRegressor
+
+model = PySRRegressor(
+ type_spec=type_spec,
+ operators={
+ 1: ["nn_relu(a) = NNValue(max.(a.data, 0.0))"],
+ 2: [
+ "nn_matmul(a, b) = NNValue(safe_matmul(a.data, b.data))",
+ "nn_add(a, b) = NNValue(safe_add(a.data, b.data))",
+ ],
+ },
+ elementwise_loss="""
+ function nn_mse(a, b)::Float64
+ valid = a.data isa Vector && b.data isa Vector && size(a.data) == size(b.data)
+ return valid ? sum(abs2, a.data .- b.data) / length(a.data) : 1.0e6
+ end
+ """,
+ niterations=100,
+ populations=4,
+ maxsize=11,
+)
+
+model.fit(X, y)
+print(model.equations_)
+```
+
+A recovered expression can have the form `nn_matmul(W2, nn_add(b, nn_relu(nn_matmul(W1, nn_add(x, c)))))`. The inner constant `c` represents the first bias through $b_1 = W_1c$, and the outer constant `b` represents the second through $b_2 = W_2b$. Each displayed constant is an `NNValue` containing its fitted scalar, vector, or matrix payload.
+
+
+
+The same custom-value interface can seed a search with `guesses`. In the vector example above,
+`Vec2` is the generated constructor from its `TypeSpec`, `add_vectors` is the declared binary
+operator, and `x1` is the first named feature column. A matching guess is
+`guesses=["add_vectors(x1, Vec2([1.0, 2.0]))"]`.
diff --git a/docs/src/migration.md b/docs/src/migration.md
new file mode 100644
index 000000000..3fa3b7437
--- /dev/null
+++ b/docs/src/migration.md
@@ -0,0 +1,207 @@
+# Migrating from v1 to v2
+
+PySR v2 upgrades the Julia backend from SymbolicRegression.jl v1.11 to v2.
+Most `PySRRegressor` constructor calls still work, but several search defaults
+changed. The old parametric-expression interface and v1 checkpoints do not
+carry over.
+
+Before upgrading, keep a copy of each run directory, including
+`hall_of_fame.csv`. Install v2 in a new Python environment if you need to keep a
+v1 environment available for old pickles or exact search reproduction.
+
+## Search defaults changed
+
+The following changes affect ordinary `PySRRegressor(...)` searches:
+
+| Setting | v1 default | v2 default |
+| --- | --- | --- |
+| `batching` | `False` | `"auto"` |
+| `batch_size` | `50` | `None` |
+| `annealing` | `False` | `True` |
+| `crossover_probability` | `0.0259` | `0.2` |
+| Mutation selection | Fixed v1 weights | New backend weights, adapted during the search |
+
+With `batching="auto"`, the backend enables batching for large datasets while
+still evaluating hall-of-fame candidates on the full dataset. With
+`batch_size=None`, it uses the full dataset for at most 1000 rows, then batches
+of 128, 256, or 512 as the dataset grows.
+
+The default mutation mix changed, including a new feature mutation that can
+replace the input referenced by a variable node. The
+`AdaptiveMutationWeightsPlugin` also updates mutation weights from observed
+success rates. A `weight_*` argument now supplies an initial weight, rather than
+a fixed weight, while that plugin is active.
+
+For a search close to the v1 defaults, set the changed values explicitly and
+omit the adaptive-mutation plugin:
+
+```python
+from pysr import AdaptiveParsimonyPlugin, PySRRegressor
+
+model = PySRRegressor(
+ batching=False,
+ batch_size=50,
+ annealing=False,
+ crossover_probability=0.0259,
+ default_plugins=[AdaptiveParsimonyPlugin()],
+ weight_add_node=2.47,
+ weight_insert_node=0.0112,
+ weight_delete_node=0.87,
+ weight_do_nothing=0.273,
+ weight_mutate_constant=0.0346,
+ weight_mutate_operator=0.293,
+ weight_mutate_feature=0.0,
+ weight_swap_operands=0.198,
+ weight_rotate_tree=4.26,
+ weight_randomize=0.000502,
+ weight_simplify=0.00209,
+ weight_optimize=0.0,
+ weight_backsolve=0.0,
+)
+```
+
+This preserves the v1 search configuration, but it does not make a run
+bit-for-bit identical. Backend implementation changes still alter random-number
+consumption and search trajectories.
+
+Two other backend changes can alter results:
+
+- Simplification now recomputes an expression's cost from scratch. Equivalent
+ expressions can therefore appear in a different hall-of-fame order.
+- Constant optimization can restart zero-valued constants, allowing trajectories
+ that v1 could not reach.
+
+Loss functions are now validated when `fit` is called. An `elementwise_loss` or
+`loss_function` with the wrong number of arguments raises a `ValueError` before
+the search starts.
+
+## Parametric expressions
+
+`ParametricExpressionSpec` and the `category=` argument to `fit` were removed.
+Use `TemplateExpressionSpec.parameters`, append the category to `X`, and index
+parameters from the template.
+
+A v1 model such as:
+
+```python
+from pysr import ParametricExpressionSpec, PySRRegressor
+
+model = PySRRegressor(
+ expression_spec=ParametricExpressionSpec(max_parameters=2),
+)
+model.fit(X, y, category=category)
+```
+
+becomes:
+
+```python
+import numpy as np
+from pysr import PySRRegressor, TemplateExpressionSpec
+
+n_categories = len(np.unique(category))
+spec = TemplateExpressionSpec(
+ combine="f(x1, x2, p1[category], p2[category])",
+ expressions=["f"],
+ variable_names=["x1", "x2", "category"],
+ parameters={"p1": n_categories, "p2": n_categories},
+)
+
+# Julia indexing starts at 1. Add one if Python categories start at 0.
+X_with_category = np.column_stack([X, category + 1])
+model = PySRRegressor(expression_spec=spec)
+model.fit(X_with_category, y)
+```
+
+The category column must also be present when calling `predict`.
+
+## Template expressions
+
+The deprecated `function_symbols=...` and positional constructor forms of
+`TemplateExpressionSpec` were replaced by explicit template fields. State the
+combining expression, inner expressions, and variable names directly:
+
+```python
+spec = TemplateExpressionSpec(
+ combine="sin(f(x1, x2)) + g(x3)^2",
+ expressions=["f", "g"],
+ variable_names=["x1", "x2", "x3"],
+)
+```
+
+This form also supports optimizable parameter vectors through
+`parameters={"p": length}`. See [Template expressions](examples.md#template-expressions)
+for complete examples.
+
+## Saved models and checkpoints
+
+v2 uses checkpoint schema 3. A v1 pickle has no compatible schema marker, so
+`pickle.load(...)` and `PySRRegressor.from_file(...)` reject it with an
+`Unsupported PySR checkpoint schema` error. Warm starts cannot cross the v1 to
+v2 boundary. Re-run the search under v2 if you need to continue evolution.
+
+For standard expressions, you can reconstruct the discovered equations from a
+copy of `hall_of_fame.csv`. Put the CSV in a run directory without the
+incompatible `checkpoint.pkl`, then provide the original operators and feature
+metadata:
+
+```python
+model = PySRRegressor.from_file(
+ run_directory="v1_csv_copy",
+ binary_operators=["+", "-", "*", "/"],
+ unary_operators=["sin", "cos"],
+ n_features_in=3,
+ feature_names_in=["x0", "x1", "x2"],
+)
+```
+
+Custom expression specifications may require a fresh search because their Julia
+objects cannot be reconstructed from the CSV alone.
+
+## v1 spellings that still work
+
+These forms emit a deprecation warning but retain their old behavior:
+
+- camelCase keyword arguments such as `fractionReplaced`, `npop`,
+ `weightMutateConstant`, and `crossoverProbability`; use snake_case instead.
+- `loss=` and `full_objective=`; use `elementwise_loss=` and `loss_function=`.
+- `multithreading=True` or `False`; use `parallelism="multithreading"` or
+ `parallelism="serial"`. Use `parallelism="multiprocessing", procs=N` for
+ process workers.
+- `.equations` and `.raw_julia_state_`; use `.equations_` and `.julia_state_`.
+- The `pysr()` function, `pysr.install()`, and `pysr.init_julia()`; use
+ `PySRRegressor`. Dependencies are installed when PySR is imported.
+
+These constructor arguments are accepted but ignored with a warning:
+
+- `julia_project` and `julia_kwargs`. Configure the backend through `juliapkg`;
+ see [Customizing the backend](backend.md).
+- `weights`, `variable_names`, and `Xresampled`. Pass them to `fit` instead.
+
+## New APIs in v2
+
+- **N-ary operators:** use
+ `operators={1: ["sin"], 2: ["+", "*"], 3: ["clamp"]}`. The legacy
+ `binary_operators` and `unary_operators` arguments still work, but cannot be
+ combined with `operators`. Constraints for n-ary operators use n-tuples, such
+ as `{"clamp": (-1, -1, 1)}`.
+- **First-class mutations, crossovers, and plugins:** configure weighted
+ instances with `mutations=` and `crossovers=`, and compose search behavior
+ with `plugins=`. See the [API reference](api.md#mutations).
+- **Custom value types:** `TypeSpec` supports expressions over strings, vectors,
+ tensors, structs, and other Julia value types. See
+ [Custom value types](examples.md#custom-value-types).
+- **Template parameters:** `TemplateExpressionSpec.parameters` supports learned
+ scalar or category-indexed parameter vectors.
+- **Equation guesses:** seed a search with `guesses=["sin(x1 * 2.1)", ...]`.
+ `fraction_replaced_guesses` controls their injection rate.
+- **Autodiff backends:** `autodiff_backend="Enzyme"` is available as an
+ experimental alternative to `"Zygote"`.
+- **Worker controls:** `worker_imports` loads Julia packages on multiprocessing
+ workers, and `worker_timeout` controls worker restart timeouts.
+- **Additional mutations:** `weight_mutate_feature` controls feature mutation.
+ Backsolve is disabled by default and can be configured with
+ `BacksolveMutation(...)` through `mutations=`.
+
+## Environment changes
+
+- The minimum versions are Python 3.10 and Julia 1.10.
diff --git a/docs/src/options.md b/docs/src/options.md
index 7a8f5c11e..3590808bc 100644
--- a/docs/src/options.md
+++ b/docs/src/options.md
@@ -245,7 +245,7 @@ The default loss is mean-square error, and weighted mean-square error.
One can pass an arbitrary Julia string to define a custom loss, using,
e.g., `elementwise_loss="myloss(x, y) = abs(x - y)^1.5"`. For more details,
see the
-[Losses](https://milescranmer.github.io/SymbolicRegression.jl/dev/losses/)
+[Losses](https://julia.pysr.ai/dev/losses)
page for SymbolicRegression.jl.
Here are some additional examples:
@@ -276,7 +276,7 @@ model = PySRRegressor(..., weights=weights, elementwise_loss="myloss(x, y, w) =
model.fit(..., weights=weights)
```
-Built-in loss (faster) (see [losses](https://ai.damtp.cam.ac.uk/symbolicregression/dev/losses/)).
+Built-in loss (faster) (see [losses](https://julia.pysr.ai/dev/losses)).
This one computes the L3 norm:
```python
diff --git a/docs/src/public/apple-touch-icon.png b/docs/src/public/apple-touch-icon.png
new file mode 100644
index 000000000..439235373
Binary files /dev/null and b/docs/src/public/apple-touch-icon.png differ
diff --git a/docs/src/public/favicon-16x16.png b/docs/src/public/favicon-16x16.png
new file mode 100644
index 000000000..d6a26a31d
Binary files /dev/null and b/docs/src/public/favicon-16x16.png differ
diff --git a/docs/src/public/favicon-32x32.png b/docs/src/public/favicon-32x32.png
new file mode 100644
index 000000000..d46b06535
Binary files /dev/null and b/docs/src/public/favicon-32x32.png differ
diff --git a/docs/src/public/favicon.ico b/docs/src/public/favicon.ico
new file mode 100644
index 000000000..ba1cdba34
Binary files /dev/null and b/docs/src/public/favicon.ico differ
diff --git a/docs/src/public/favicon.png b/docs/src/public/favicon.png
index 5fcc0ab9f..69498a477 100644
Binary files a/docs/src/public/favicon.png and b/docs/src/public/favicon.png differ
diff --git a/docs/src/slurm.md b/docs/src/slurm.md
new file mode 100644
index 000000000..ec1c395bf
--- /dev/null
+++ b/docs/src/slurm.md
@@ -0,0 +1,50 @@
+# Slurm (Multi-node)
+
+PySR supports running across multiple nodes in an existing Slurm allocation with
+`cluster_manager="slurm"`. Request the resources with `sbatch` or `salloc`, then
+PySR will launch Julia workers across the allocation using
+`SlurmClusterManager.jl`.
+
+For example, this job requests three workers on each of two nodes:
+
+```bash
+#!/bin/bash
+#SBATCH --job-name=pysr
+#SBATCH --partition=normal
+#SBATCH --nodes=2
+#SBATCH --ntasks-per-node=3
+#SBATCH --time=01:00:00
+
+set -euo pipefail
+python pysr_script.py
+```
+
+The corresponding `pysr_script.py` is:
+
+```python
+import numpy as np
+from pysr import PySRRegressor
+
+X = np.random.RandomState(0).randn(1000, 2)
+y = X[:, 0] + 2 * X[:, 1]
+
+model = PySRRegressor(
+ niterations=200,
+ populations=2,
+ parallelism="multiprocessing",
+ cluster_manager="slurm",
+ procs=6,
+)
+model.fit(X, y)
+print(model)
+```
+
+Submit the job with:
+
+```bash
+sbatch pysr_job.sh
+```
+
+`procs` must equal the allocation's total task count, given by `--ntasks` or
+`--nodes` multiplied by `--ntasks-per-node`. Run the Python script once inside
+the allocation; do not wrap it in `srun`.
diff --git a/docs/src/tuning.md b/docs/src/tuning.md
index 443789e49..ec995d29b 100644
--- a/docs/src/tuning.md
+++ b/docs/src/tuning.md
@@ -6,11 +6,11 @@ First, my general tips would be to avoid using redundant operators, like how `po
When running PySR, I usually do the following:
-I run from IPython (Jupyter Notebooks don't work as well[^1]) on the head node of a slurm cluster. Passing `cluster_manager="slurm"` will make PySR set up a run over the entire allocation. I set `procs` equal to the total number of cores over my entire allocation.
+I run from IPython (Jupyter Notebooks don't work as well[^1]) on the head node of a slurm cluster. Passing `cluster_manager="slurm"` will make PySR set up a run over the entire allocation. I set `procs` equal to the total number of tasks across my entire allocation (see the [Slurm page](slurm.md) for a complete multi-node example).
-I use the [tensorboard feature](https://ai.damtp.cam.ac.uk/pysr/examples/#12-using-tensorboard-for-logging) for experiment tracking.
+I use the [tensorboard feature](/examples/instrumentation#using-tensorboard-for-logging) for experiment tracking.
-[^1]: Jupyter Notebooks are supported by PySR, but miss out on some useful features available in IPython and Python: the progress bar, and early stopping with "q". In Jupyter you cannot interrupt a search once it has started; you have to restart the kernel. See [this issue](https://github.com/astroautomata/PySR/issues/260) for updates.
+[^1]: Jupyter Notebooks are supported by PySR, but miss out on some useful features available in IPython and Python: the progress bar, and early stopping with "q". Interrupting a search in Jupyter returns the equations found so far.
1. I start by using the default parameters.
2. I use only the operators I think it needs and no more.
diff --git a/docs/test_deploy_versions.jl b/docs/test_deploy_versions.jl
new file mode 100644
index 000000000..886bd939d
--- /dev/null
+++ b/docs/test_deploy_versions.jl
@@ -0,0 +1,163 @@
+using Test
+
+include(joinpath(@__DIR__, "deploy_versions.jl"))
+
+@testset "stable ownership" begin
+ tags = [v"2.1.0", v"2.0.0", v"1.5.10"]
+
+ @test claims_stable("v2.1.0"; all_tagged_versions = tags, log = false)
+ @test !claims_stable("v2.0.0"; all_tagged_versions = tags, log = false)
+ # determine_bases would report this alpha as owning v2.0.0
+ @test !claims_stable("v2.0.0a2"; all_tagged_versions = tags, log = false)
+ @test !claims_stable("dev"; all_tagged_versions = tags, log = false)
+ @test !claims_stable("previews/PR123"; all_tagged_versions = tags, log = false)
+end
+
+@testset "stable deployment root" begin
+ @test stable_deploy_root("/tmp/gh-pages/stable") == "/tmp/gh-pages"
+ @test stable_deploy_root("/tmp/gh-pages/stable/") == "/tmp/gh-pages"
+end
+
+@testset "stable redirects" begin
+ pages = ["index", "examples", "404", "api"]
+ files = stable_redirect_files(pages, "v2.2.1")
+
+ @test files["api.html"] == """
+
+"""
+ @test files["api/index.html"] == """
+
+"""
+ @test files["stable/api.html"] == """
+
+"""
+ @test files["stable/index.html"] == """
+
+"""
+
+ @test sort(collect(keys(files))) == [
+ "api.html",
+ "api/index.html",
+ "examples.html",
+ "examples/index.html",
+ "stable/api.html",
+ "stable/examples.html",
+ "stable/index.html",
+ ]
+ @test collect(files) == collect(stable_redirect_files(reverse(pages), "v2.2.1"))
+ root_stubs = filter(((path, _),) -> !startswith(path, "stable/"), files)
+ @test all(!occursin(r"v\d", content) for content in values(root_stubs))
+end
+
+@testset "stale redirects" begin
+ files = stable_redirect_files(["api"], "v2.2.1")
+ root_stubs = filter(((path, _),) -> !startswith(path, "stable/"), files)
+
+ root = mktempdir()
+ for (path, content) in root_stubs
+ mkpath(joinpath(root, Base.dirname(path)))
+ write(joinpath(root, path), content)
+ end
+ # A page dropped since the previous release, in both of its stub forms.
+ for (path, content) in stable_redirect_files(["gone"], "v2.2.0")
+ startswith(path, "stable/") && continue
+ mkpath(joinpath(root, Base.dirname(path)))
+ write(joinpath(root, path), content)
+ end
+ # Everything a deployment leaves at the root besides our stubs.
+ write(joinpath(root, "index.html"), "\n")
+ write(joinpath(root, "versions.js"), "var DOC_VERSIONS = [];\n")
+ write(joinpath(root, "CNAME"), "pysr.ai\n")
+ for version in ("v2.2.1", "dev", "stable")
+ mkpath(joinpath(root, version))
+ write(joinpath(root, version, "index.html"), "\n")
+ end
+ write(joinpath(root, "stable", "index.html"), files["stable/index.html"])
+
+ @test stale_root_stubs(root, keys(root_stubs)) == ["gone.html", "gone/index.html"]
+ @test read(joinpath(root, "CNAME"), String) == "pysr.ai\n"
+ @test isempty(stale_root_stubs(root, keys(stable_redirect_files(["api", "gone"], "v2.2.1"))))
+end
+
+@testset "nested redirects" begin
+ files = stable_redirect_files(["api", "examples/objectives", "guide/index"], "v2.3.0")
+
+ @test files["examples/objectives.html"] == """$(STUB_MARKER)
+
+"""
+ @test files["examples/objectives/index.html"] == """$(STUB_MARKER)
+
+"""
+ @test files["stable/examples/objectives.html"] == """$(STUB_MARKER)
+
+"""
+ # A page that is already a directory index gets no `/index.html` of its own.
+ @test files["guide/index.html"] == """$(STUB_MARKER)
+
+"""
+ @test !haskey(files, "guide/index/index.html")
+ @test sort(collect(keys(files))) == [
+ "api.html",
+ "api/index.html",
+ "examples/objectives.html",
+ "examples/objectives/index.html",
+ "guide/index.html",
+ "stable/api.html",
+ "stable/examples/objectives.html",
+ "stable/guide/index.html",
+ "stable/index.html",
+ ]
+ root_stubs = filter(((path, _),) -> !startswith(path, "stable/"), files)
+ @test all(!occursin(r"v\d", content) for content in values(root_stubs))
+ # A flat page and a directory index cannot both own the same stub path.
+ @test_throws ErrorException stable_redirect_files(["guide", "guide/index"], "v2.3.0")
+end
+
+@testset "built pages" begin
+ dist = mktempdir()
+ mkpath(joinpath(dist, "examples"))
+ mkpath(joinpath(dist, "assets"))
+ for path in ("index.html", "404.html", "api.html", "examples/objectives.html")
+ write(joinpath(dist, path), "\n")
+ end
+ write(joinpath(dist, "assets", "app.js"), "")
+ write(joinpath(dist, "hashmap.json"), "{}")
+
+ @test built_pages(dist) == ["404", "api", "examples/objectives", "index"]
+ @test sort(collect(keys(stable_redirect_files(built_pages(dist), "v2.3.0")))) == [
+ "api.html",
+ "api/index.html",
+ "examples/objectives.html",
+ "examples/objectives/index.html",
+ "stable/api.html",
+ "stable/examples/objectives.html",
+ "stable/index.html",
+ ]
+end
+
+@testset "nested stale redirects" begin
+ current = stable_redirect_files(["api"], "v2.3.0")
+ root_stubs = filter(((path, _),) -> !startswith(path, "stable/"), current)
+
+ root = mktempdir()
+ for (path, content) in root_stubs
+ mkpath(joinpath(root, Base.dirname(path)))
+ write(joinpath(root, path), content)
+ end
+ # Stubs for a nested page that has since been removed or renamed.
+ for (path, content) in stable_redirect_files(["examples/objectives"], "v2.2.1")
+ startswith(path, "stable/") && continue
+ mkpath(joinpath(root, Base.dirname(path)))
+ write(joinpath(root, path), content)
+ end
+ # Real builds, whose pages carry no marker and so are never candidates.
+ for version in ("v2.2.1", "dev", "previews/PR1", "stable")
+ mkpath(joinpath(root, version, "examples"))
+ write(joinpath(root, version, "examples", "objectives.html"), "\n")
+ end
+ write(joinpath(root, "previews", "PR1", "stale.html"), STUB_MARKER * "\n")
+ write(joinpath(root, "stable", "stale.html"), STUB_MARKER * "\n")
+
+ @test stale_root_stubs(root, keys(root_stubs)) ==
+ ["examples/objectives.html", "examples/objectives/index.html"]
+end
diff --git a/environment.yml b/environment.yml
index 6ef91fbe5..587f9b410 100644
--- a/environment.yml
+++ b/environment.yml
@@ -2,11 +2,11 @@ name: test
channels:
- conda-forge
dependencies:
- - python>=3.9
- - sympy>=1.0.0,<2.0.0
+ - python>=3.10
+ - sympy>=1.14.0,<2.0.0
- pandas>=0.21.0,<4.0.0
- numpy>=1.13.0,<3.0.0
- scikit-learn>=1.0.0,<2.0.0
- - pyjuliacall>=0.9.28,<0.9.29
+ - pyjuliacall>=0.9.29,<0.9.36
- click>=7.0.0,<9.0.0
- typing-extensions>=4.0.0,<5.0.0
diff --git a/examples/adaptive_mutation_weights.py b/examples/adaptive_mutation_weights.py
new file mode 100644
index 000000000..2fd92828a
--- /dev/null
+++ b/examples/adaptive_mutation_weights.py
@@ -0,0 +1,305 @@
+"""Watch the search re-allocate its own mutation budget.
+
+`AdaptiveMutationWeightsPlugin` is on by default: per mutation kind it counts attempts
+and strictly improving children, and scales that kind's weight by the learned ratio.
+SymbolicRegression.jl logs none of those counters, so this example puts the shipped
+plugin inside a plugin of its own that forwards every hook to it and reads its state
+on the way out, then prints where the budget ended up.
+"""
+
+from dataclasses import dataclass, field
+
+import numpy as np
+
+from pysr import PySRRegressor
+from pysr.julia_import import AnyValue, jl
+from pysr.plugins import (
+ AbstractPlugin,
+ AdaptiveMutationWeightsPlugin,
+ AdaptiveParsimonyPlugin,
+ SimulatedAnnealingPlugin,
+)
+
+_rng = np.random.default_rng(20260817)
+X = _rng.uniform(-3.0, 3.0, size=(200, 5))
+y = 2.5382 * np.cos(X[:, 3]) + X[:, 0] ** 2 - 0.5
+VARIABLE_NAMES = None
+
+# `DocsTracedAdaptive` holds a real `AdaptiveMutationWeightsPlugin` and forwards
+# `init_plugin_state`, `on_mutation_end!` and `condition_mutation_weights!` to it
+# unchanged, so the adaptation arithmetic during the run is the library's own. The
+# wrapper only counts draws per mutation kind and snapshots the plugin's own
+# `multipliers` vector as it moves.
+_TRACER_JL = """
+if !isdefined(Main, :DocsTracedAdaptive)
+
+const DocsAMW = SymbolicRegression.AdaptiveMutationWeightsModule
+
+struct DocsTracedAdaptive <: SymbolicRegression.AbstractPlugin
+ inner::DocsAMW.AdaptiveMutationWeightsPlugin
+ stride::Int
+end
+
+DocsTracedAdaptive(; stride::Int=1000, kws...) =
+ DocsTracedAdaptive(DocsAMW.AdaptiveMutationWeightsPlugin(; kws...), stride)
+
+mutable struct DocsTracedState
+ inner::DocsAMW.AdaptiveMutationWeightsState
+ draws::Vector{Float64}
+ seen::Int
+ snapshots::Vector{Tuple{Int,Vector{Float64}}}
+ registered::Bool
+end
+
+const DOCS_TRACED_STATES = DocsTracedState[]
+
+function SymbolicRegression.init_plugin_state(p::DocsTracedAdaptive, options, dataset)
+ inner = SymbolicRegression.init_plugin_state(p.inner, options, dataset)
+ n = length(options.mutations)
+ return DocsTracedState(inner, zeros(n), 0, Tuple{Int,Vector{Float64}}[], false)
+end
+
+function SymbolicRegression.on_mutation_end!(
+ s::DocsTracedState,
+ p::DocsTracedAdaptive,
+ mutation::SymbolicRegression.AbstractMutation,
+ event::SymbolicRegression.MutationEvent,
+ dataset,
+ options::SymbolicRegression.AbstractOptions,
+)
+ if !s.registered
+ push!(DOCS_TRACED_STATES, s)
+ s.registered = true
+ end
+ s.draws[event.mutation_idx] += 1.0
+ s.seen += 1
+ SymbolicRegression.on_mutation_end!(s.inner, p.inner, mutation, event, dataset, options)
+ if s.seen % p.stride == 0
+ push!(s.snapshots, (s.seen, copy(s.inner.multipliers)))
+ end
+ return nothing
+end
+
+function SymbolicRegression.condition_mutation_weights!(
+ weights::AbstractVector,
+ s::DocsTracedState,
+ p::DocsTracedAdaptive,
+ member,
+ options::SymbolicRegression.AbstractOptions,
+ curmaxsize,
+ nfeatures,
+)
+ return SymbolicRegression.condition_mutation_weights!(
+ weights, s.inner, p.inner, member, options, curmaxsize, nfeatures
+ )
+end
+
+end
+"""
+
+
+@dataclass(frozen=True)
+class TracedAdaptiveMutationWeights(AbstractPlugin):
+ """`AdaptiveMutationWeightsPlugin`, wrapped so its learned state can be read."""
+
+ stride: int = 1000
+ inner: AdaptiveMutationWeightsPlugin = field(
+ default_factory=AdaptiveMutationWeightsPlugin
+ )
+
+ def julia_plugin(self) -> AnyValue:
+ jl.seval(_TRACER_JL)
+ # PySR calls this once per `fit`, while building that fit's options.
+ jl.seval("empty!(DOCS_TRACED_STATES)")
+ return jl.DocsTracedAdaptive(
+ stride=self.stride,
+ smoothing=self.inner.smoothing,
+ floor=self.inner.floor,
+ reward=jl.Symbol(self.inner.reward),
+ )
+
+
+TRACER = TracedAdaptiveMutationWeights()
+
+MODEL_KWARGS = dict(
+ binary_operators=["+", "-", "*", "/"],
+ unary_operators=["cos", "exp"],
+ maxsize=14,
+ niterations=300,
+ # The plugin keeps one set of counters per population, so a single population is
+ # what makes the printed numbers a run rather than an average over runs.
+ populations=1,
+ # The adaptive plugin is a *default* plugin, so it has to be replaced rather than
+ # added: an extra plugin of a new type leaves the shipped one in place and the
+ # search adapts twice. `alpha=3.17` is PySR's own default, asserted in `main`.
+ plugins=[],
+ default_plugins=[
+ SimulatedAnnealingPlugin(alpha=3.17),
+ AdaptiveParsimonyPlugin(),
+ TRACER,
+ ],
+ verbosity=0,
+)
+
+
+def check(model):
+ """The target recovered to numerical precision."""
+ return bool(model.equations_["loss"].min() < 1e-12)
+
+
+def _states():
+ return int(jl.seval("length(DOCS_TRACED_STATES)"))
+
+
+def _pooled(expr, n):
+ """Sum a per-state vector over every population's state."""
+ total = [0.0] * n
+ for s in range(1, _states() + 1):
+ for i, v in enumerate(jl.seval(f"DOCS_TRACED_STATES[{s}].{expr}")):
+ total[i] += float(v)
+ return total
+
+
+def mutation_statistics(model):
+ """Per-kind counters read out of the plugin's own state after a fit."""
+ jl.DOCS_TRACED_OPTIONS = model.julia_options_
+ names = [
+ str(v)
+ for v in jl.seval(
+ "[string(typeof(first(p))) for p in DOCS_TRACED_OPTIONS.mutations]"
+ )
+ ]
+ base = [
+ float(v)
+ for v in jl.seval("[Float64(last(p)) for p in DOCS_TRACED_OPTIONS.mutations]")
+ ]
+ active = [bool(v) for v in jl.seval("DOCS_TRACED_STATES[1].inner.active")]
+ n = len(names)
+ n_states = _states()
+
+ draws = _pooled("draws", n)
+ attempts = _pooled("inner.attempts", n)
+ successes = _pooled("inner.successes", n)
+ multipliers = [m / n_states for m in _pooled("inner.multipliers", n)]
+ strength = float(
+ jl.seval("DocsAMW.AdaptiveMutationWeightsPlugin().adaptation_strength")
+ )
+
+ # The kinds the run actually drew from: nonzero base weight, and legal for a plain
+ # expression. Effective weight is base * multiplier ** adaptation_strength.
+ drawn = [i for i in range(n) if draws[i] > 0]
+ eff = {i: base[i] * multipliers[i] ** strength for i in drawn}
+ base_total = sum(base[i] for i in drawn)
+ eff_total = sum(eff.values())
+ draw_total = sum(draws[i] for i in drawn)
+
+ rows = [
+ {
+ "mutation": names[i],
+ "index": i + 1,
+ "adapts": active[i],
+ "base_weight": base[i],
+ "draws": int(draws[i]),
+ "attempts": int(attempts[i]),
+ "successes": int(successes[i]),
+ # The quantity the plugin scores with: a Laplace-smoothed improvement rate.
+ "rate": (successes[i] + 1.0) / (attempts[i] + 2.0) if active[i] else None,
+ "multiplier": multipliers[i],
+ "share_shipped": base[i] / base_total,
+ "share_learned": eff[i] / eff_total,
+ "share_drawn": draws[i] / draw_total,
+ }
+ for i in drawn
+ ]
+ rows.sort(key=lambda r: -r["multiplier"])
+ meta = {
+ "n_states": n_states,
+ "n_mutations": int(sum(draws)),
+ "adaptation_strength": strength,
+ "plugins": [
+ str(v)
+ for v in jl.seval(
+ "[string(typeof(p)) for p in DOCS_TRACED_OPTIONS.plugins]"
+ )
+ ],
+ }
+ return rows, meta
+
+
+def multiplier_trajectory(rows, n_points=6):
+ """The learned multipliers at geometrically spaced points through the run.
+
+ Geometric rather than even: most of the re-allocation happens in the first few
+ thousand mutations, which even spacing would step straight over.
+ """
+ n_states = _states()
+ n_snaps = min(
+ int(jl.seval(f"length(DOCS_TRACED_STATES[{s}].snapshots)"))
+ for s in range(1, n_states + 1)
+ )
+ if n_snaps == 0:
+ return []
+ picks = sorted(
+ {1} | {max(1, round(n_snaps ** (k / n_points))) for k in range(1, n_points + 1)}
+ )
+ out = []
+ for p in picks:
+ seen = 0
+ totals = None
+ for s in range(1, n_states + 1):
+ snap = f"DOCS_TRACED_STATES[{s}].snapshots[{p}]"
+ seen += int(jl.seval(f"{snap}[1]"))
+ values = [float(v) for v in jl.seval(f"{snap}[2]")]
+ totals = (
+ values if totals is None else [a + b for a, b in zip(totals, values)]
+ )
+ out.append((seen, [totals[r["index"] - 1] / n_states for r in rows]))
+ return out
+
+
+def report(rows, meta):
+ print(
+ f"\n{meta['n_mutations']} mutations, {meta['n_states']} population state(s), adaptation_strength={meta['adaptation_strength']}, plugins={meta['plugins']}"
+ )
+
+ print(
+ f"\n{'mutation':<23}{'draws':>8}{'improves':>10}{'multiplier':>12}{'share: shipped':>16}{'learned':>10}{'drawn':>9}"
+ )
+ for r in rows:
+ rate = " -" if r["rate"] is None else f"{r['rate']:10.4f}"
+ print(
+ f"{r['mutation']:<23}{r['draws']:>8}{rate}{r['multiplier']:>12.3f}{r['share_shipped']:>16.3%}{r['share_learned']:>10.3%}{r['share_drawn']:>9.3%}"
+ )
+ print("\nshipped and learned are weight shares; drawn is what the search sampled,")
+ print("which also reflects the engine zeroing kinds illegal for a given member.")
+
+ trajectory = multiplier_trajectory(rows)
+ if trajectory:
+ print("\nlearned multiplier through the run, by mutations elapsed:")
+ print(f"{'mutation':<23}" + "".join(f"{seen:>10}" for seen, _ in trajectory))
+ for j, r in enumerate(rows):
+ print(
+ f"{r['mutation']:<23}"
+ + "".join(f"{values[j]:>10.3f}" for _, values in trajectory)
+ )
+
+
+def main():
+ assert MODEL_KWARGS["default_plugins"][0].alpha == PySRRegressor().alpha
+
+ model = PySRRegressor(**MODEL_KWARGS, random_state=0)
+ model.fit(X, y, variable_names=VARIABLE_NAMES)
+ print(model.equations_[["complexity", "loss", "equation"]].to_string(index=False))
+
+ rows, meta = mutation_statistics(model)
+ assert meta["plugins"] == [
+ "SimulatedAnnealingPlugin",
+ "AdaptiveParsimonyPlugin",
+ "DocsTracedAdaptive",
+ ], meta["plugins"]
+ report(rows, meta)
+ print(f"\nrecovered exactly: {check(model)}")
+
+
+if __name__ == "__main__":
+ main()
diff --git a/examples/affine_cipher.py b/examples/affine_cipher.py
new file mode 100644
index 000000000..e17339e3d
--- /dev/null
+++ b/examples/affine_cipher.py
@@ -0,0 +1,69 @@
+"""Break an affine cipher with symbolic regression over a custom letter type.
+
+The cipher is E(p) = (5p + 8) mod 26 on the 26-letter alphabet. The search sees
+only pairs of (ciphertext letter, plaintext letter) and has to discover the
+decryption law. Letters are a custom type rather than floats, so the operators are
+modular letter arithmetic and the constants in the printed expressions are letters.
+"""
+
+import numpy as np
+
+from pysr import PySRRegressor, TypeSpec
+
+ALPHABET = "ABCDEFGHIJKLMNOPQRSTUVWXYZ"
+PLAIN = "THELAWSOFPHYSICSAREWRITTENINTHELANGUAGEOFMATHEMATICS"
+CIPHER = "".join(ALPHABET[(5 * ALPHABET.index(c) + 8) % 26] for c in PLAIN)
+
+X = np.array([[ALPHABET.index(c)] for c in CIPHER], dtype=object)
+y = np.array([ALPHABET.index(c) for c in PLAIN], dtype=object)
+
+# A letter is an integer code in 0:25. Mutation walks one letter up or down, and
+# the string form prints the letter itself, so a constant reads as `V`, not 21.
+SPEC = TypeSpec(
+ "Letter",
+ fields={"code": "Int"},
+ sample="rng -> Letter(rand(rng, 0:25))",
+ mutate="(rng, value, temperature) -> Letter(mod(value.code + rand(rng, (-1, 1)), 26))",
+ string="value -> string(Char('A' + value.code))",
+)
+
+# Distance in letter codes between the prediction and the target, which PySR averages
+# over the rows. It is zero exactly when all 52 letters are right, and a hit-or-miss
+# count would score a candidate one letter off the same as one twenty off.
+LETTER_LOSS = "letter_loss(prediction::Letter, target::Letter)::Float64 = abs(prediction.code - target.code)"
+
+MODEL_KWARGS = dict(
+ type_spec=SPEC,
+ # Modular arithmetic on letters: every operator stays inside the alphabet.
+ operators={
+ 2: [
+ "shift(a::Letter, b::Letter) = Letter(mod(a.code + b.code, 26))",
+ "unshift(a::Letter, b::Letter) = Letter(mod(a.code - b.code, 26))",
+ "mix(a::Letter, b::Letter) = Letter(mod(a.code * b.code, 26))",
+ ]
+ },
+ elementwise_loss=LETTER_LOSS,
+ niterations=20,
+)
+
+
+def decode(model):
+ return "".join(ALPHABET[int(v) % 26] for v in model.predict(X))
+
+
+def check(model):
+ """The whole message comes back, letter for letter."""
+ return decode(model) == PLAIN
+
+
+def main():
+ model = PySRRegressor(**MODEL_KWARGS, random_state=0)
+ model.fit(X, y)
+ print(model.equations_[["complexity", "loss", "equation"]].to_string(index=False))
+ print(f"\ndecoded: {decode(model)}")
+ print(f"plaintext: {PLAIN}")
+ print(f"exact: {check(model)}")
+
+
+if __name__ == "__main__":
+ main()
diff --git a/examples/affine_freemod.py b/examples/affine_freemod.py
new file mode 100644
index 000000000..1e9d052d2
--- /dev/null
+++ b/examples/affine_freemod.py
@@ -0,0 +1,90 @@
+"""Affine cipher with the modulus withheld: 26 has to be learned, not declared.
+
+The data is 52 (ciphertext, plaintext) letter-code pairs from E(p) = (5p + 8) mod 26.
+The search sees only the codes. The value type is a plain integer in 0:51 and the three
+operators are integer `+`, `*` and `mod` with nothing folded in, so the decryption law
+`mod(21*x0 + 14, 26)` has to be written out of three learned constants, the 26 among
+them.
+"""
+
+import numpy as np
+
+from pysr import PySRRegressor, TypeSpec
+
+ALPHABET = "ABCDEFGHIJKLMNOPQRSTUVWXYZ"
+A, B = 5, 8
+PLAIN = "THELAWSOFPHYSICSAREWRITTENINTHELANGUAGEOFMATHEMATICS"
+CIPHER = "".join(ALPHABET[(A * ALPHABET.index(c) + B) % 26] for c in PLAIN)
+
+X = np.array([[ALPHABET.index(c)] for c in CIPHER], dtype=object)
+y = np.array([ALPHABET.index(c) for c in PLAIN], dtype=object)
+
+# The sampled constant range. 26 must not be the only value that fits in it, so the
+# range runs to twice the alphabet and every modulus up to 51 is reachable.
+CMAX = 51
+
+# Constant-proposal jump rate: the share of proposals that resample the whole range
+# instead of stepping by one. One of the three constants is a modulus, and 25 or 27 in
+# that slot scrambles the output as badly as any other wrong value, so a walk alone
+# cannot find it.
+JUMP = 0.35
+
+SPEC = TypeSpec(
+ "Num",
+ fields={"n": "Int"},
+ sample=f"rng -> Num(rand(rng, 0:{CMAX}))",
+ mutate=f"(rng, value, temperature) -> Num(rand(rng) < {JUMP} ? rand(rng, 0:{CMAX}) : clamp(value.n + rand(rng, (-1, 1)), 0, {CMAX}))",
+ string="value -> string(value.n)",
+)
+
+# Integer arithmetic, no alphabet size anywhere. `mod` needs a total extension because
+# a subtree can evaluate to zero; a zero modulus returns the dividend, which keeps the
+# operator total without making any candidate look better than it is.
+OPERATORS = {
+ 2: [
+ "add(a::Num, b::Num) = Num(a.n + b.n)",
+ "mul(a::Num, b::Num) = Num(a.n * b.n)",
+ "modulo(a::Num, b::Num) = Num(b.n == 0 ? a.n : mod(a.n, b.n))",
+ ]
+}
+
+# Absolute distance between the candidate's number and the target letter code, which
+# PySR averages over the rows. It is zero exactly when every one of the 52 letters is
+# right, and a hit-or-miss score would give a candidate that is one off the same loss as
+# one that is twenty off, leaving the constant search nothing to climb.
+NUM_LOSS = (
+ "num_loss(prediction::Num, target::Num)::Float64 = abs(prediction.n - target.n)"
+)
+
+MODEL_KWARGS = dict(
+ type_spec=SPEC,
+ operators=OPERATORS,
+ elementwise_loss=NUM_LOSS,
+ niterations=1000,
+ maxsize=12,
+ should_optimize_constants=False,
+ verbosity=0,
+)
+
+
+def decode(model):
+ return "".join(ALPHABET[int(v) % 26] for v in model.predict(X))
+
+
+def check(model):
+ """The whole message comes back, letter for letter."""
+ return decode(model) == PLAIN
+
+
+def main():
+ model = PySRRegressor(**MODEL_KWARGS, random_state=0)
+ model.fit(X, y)
+
+ print(model.equations_[["complexity", "loss", "equation"]].to_string(index=False))
+ print(f"\ndecoded: {decode(model)}")
+ print(f"plaintext: {PLAIN}")
+ print(f"exact: {check(model)}")
+
+
+if __name__ == "__main__":
+ main()
diff --git a/examples/agent_loop_guesses.py b/examples/agent_loop_guesses.py
new file mode 100644
index 000000000..4ca34f51c
--- /dev/null
+++ b/examples/agent_loop_guesses.py
@@ -0,0 +1,91 @@
+"""Close an agent loop around PySR: read the front, write the next equation, search again.
+
+The data is 160 measurements of kinetic energy against mass and velocity in natural units,
+at velocities high enough that the Newtonian form is only an approximation. A cold search
+climbs a ladder of polynomials in velocity and stops three to four orders of magnitude
+above the measurement noise. Reading the front says what is missing: the relativistic
+law, passed back in through `guesses=`.
+"""
+
+import numpy as np
+
+from pysr import PySRRegressor
+
+rng = np.random.default_rng(42)
+mass = rng.uniform(0.5, 3.0, 160)
+velocity = rng.uniform(0.08, 0.88, 160)
+kinetic_energy_exact = mass * (1 / np.sqrt(1 - velocity**2) - 1)
+NOISE_SIGMA = 0.0015 * float(np.median(kinetic_energy_exact))
+kinetic_energy = kinetic_energy_exact + rng.normal(0.0, NOISE_SIGMA, mass.size)
+NOISE_FLOOR = NOISE_SIGMA**2
+
+X = np.stack([mass, velocity], axis=1)
+y = kinetic_energy
+VARIABLE_NAMES = ["mass", "velocity"]
+
+MODEL_KWARGS = dict(
+ operators={1: ["sqrt"], 2: ["+", "-", "*", "/"]},
+ niterations=10,
+ verbosity=0,
+)
+
+# The filmed session handed the front to a coding agent, which replied with this. Each
+# entry is one round's proposal, so the loop stops when the list runs out.
+PROPOSALS = ["(mass / sqrt(1 - (velocity * velocity))) - mass"]
+
+
+def agent(front, tried):
+ """Stand in for the model call: show the front, return the next expression, or None.
+
+ A real loop prompts a model right here with the front printed below and the list of
+ expressions already tried, and parses one expression out of the reply.
+ """
+ print("front handed to the agent:")
+ print(
+ front[["complexity", "loss", "equation"]].to_string(index=False),
+ "\n",
+ flush=True,
+ )
+ return PROPOSALS[len(tried)] if len(tried) < len(PROPOSALS) else None
+
+
+def check(model):
+ """The search reached the measurement noise floor."""
+ return bool(model.equations_["loss"].min() <= 2 * NOISE_FLOOR)
+
+
+def main(seed=0):
+ guesses, tried, rounds = None, [], []
+ while True:
+ model = PySRRegressor(**MODEL_KWARGS, guesses=guesses, random_state=seed)
+ model.fit(X, y, variable_names=VARIABLE_NAMES)
+ front = model.equations_
+ best = front.loc[front["loss"].idxmin()]
+ print(
+ f"round {len(rounds)} guess: {guesses[0] if guesses else 'none (cold start)'}"
+ )
+ print(f" loss {best['loss']:.4g} at complexity {best['complexity']}")
+ print(f" {best['equation']}\n", flush=True)
+ rounds.append(
+ dict(
+ guess=guesses[0] if guesses else None,
+ loss=float(best["loss"]),
+ complexity=int(best["complexity"]),
+ equation=str(best["equation"]),
+ )
+ )
+ if check(model):
+ break
+ proposal = agent(front, tried)
+ if proposal is None:
+ break
+ tried.append(proposal)
+ guesses = [proposal] # single output: a plain list of expression strings
+
+ print(front[["complexity", "loss", "equation"]].to_string(index=False))
+ print(f"\nnoise floor {NOISE_FLOOR:.4g}; reached: {check(model)}")
+ return rounds
+
+
+if __name__ == "__main__":
+ main()
diff --git a/examples/any_arity.py b/examples/any_arity.py
new file mode 100644
index 000000000..af287ae59
--- /dev/null
+++ b/examples/any_arity.py
@@ -0,0 +1,101 @@
+"""Rediscover a clipped envelope written with a four-argument minimum.
+
+Operators are keyed by arity, so an operator of any arity can be supplied: here a
+three-argument conditional and a four-argument minimum sit side by side in the
+same set. The target `min4(x + 2, 2 - x, sin(3x) + 1.2, 0.8)` folds a rising
+ramp, a falling ramp, a shifted sinusoid and a flat cap together in one call,
+which binary `min` would need three nested calls to write.
+
+Recovery is judged on 4001 points the search never saw, not on the 400 it fitted.
+"""
+
+import numpy as np
+import sympy
+
+from pysr import PySRRegressor
+
+DOMAIN = (-4.0, 4.0)
+
+
+def envelope(x):
+ return np.minimum.reduce(
+ [x + 2.0, 2.0 - x, np.sin(3.0 * x) + 1.2, np.full_like(x, 0.8)]
+ )
+
+
+_x = np.sort(np.random.default_rng(0).uniform(*DOMAIN, 400))
+X = _x.reshape(-1, 1).astype(np.float32)
+y = envelope(_x).astype(np.float32)
+VARIABLE_NAMES = ["x"]
+
+_unseen = np.linspace(*DOMAIN, 4001)
+X_UNSEEN = _unseen.reshape(-1, 1).astype(np.float32)
+Y_UNSEEN = envelope(_unseen)
+
+MODEL_KWARGS = dict(
+ operators={
+ 1: ["sin", "cos"],
+ 2: ["+", "-", "*"],
+ 3: ["ifelse(t, a, b) = t > 0 ? a : b"],
+ 4: ["min4(a, b, c, d) = min(a, b, c, d)"],
+ },
+ # A custom operator is rejected at fit time unless it has a SymPy image.
+ extra_sympy_mappings={
+ "ifelse": lambda t, a, b: sympy.Piecewise((a, t > 0), (b, True)),
+ "min4": lambda a, b, c, d: sympy.Min(a, b, c, d),
+ },
+ niterations=500,
+ verbosity=0,
+)
+
+TOLERANCE = 1e-4
+
+
+def unseen_errors(model):
+ """Largest error of each front member on the points the search never saw."""
+ return {
+ i: float(
+ np.abs(
+ np.asarray(model.predict(X_UNSEEN, index=i), np.float64) - Y_UNSEEN
+ ).max()
+ )
+ for i in model.equations_.index
+ }
+
+
+def check(model):
+ """A front member reproduces the law off the fitted samples, not just on them."""
+ return bool(any(e < TOLERANCE for e in unseen_errors(model).values()))
+
+
+def main():
+ model = PySRRegressor(**MODEL_KWARGS, random_state=0)
+ model.fit(X, y, variable_names=VARIABLE_NAMES)
+
+ front = model.equations_[["complexity", "loss", "equation"]]
+ errors = unseen_errors(model)
+ print(
+ front.assign(unseen_error=[errors[i] for i in front.index]).to_string(
+ index=False
+ )
+ )
+
+ uses = front["equation"].str.contains
+ print(
+ f"\nfront members using the arity-3 ifelse: {int(uses('ifelse').sum())}, using the arity-4 min4: {int(uses('min4').sum())}"
+ )
+
+ print("true law: min4(x + 2, 2 - x, sin(3x) + 1.2, 0.8)")
+ recovered = [i for i, e in errors.items() if e < TOLERANCE]
+ if recovered:
+ best = min(recovered, key=lambda i: front.complexity[i])
+ print(
+ f"recovered at complexity {int(front.complexity[best])}, max error {errors[best]:.3g} over 4001 unseen points:"
+ )
+ print(f" {front.equation[best]}")
+ else:
+ print("no front member agrees with the law off the fitted samples")
+
+
+if __name__ == "__main__":
+ main()
diff --git a/examples/automatic_batching.py b/examples/automatic_batching.py
new file mode 100644
index 000000000..30f7afa6b
--- /dev/null
+++ b/examples/automatic_batching.py
@@ -0,0 +1,46 @@
+"""Recover y = 2.5382*cos(x3) + x0^2 - 0.5 from 20,000 rows.
+
+PySR defaults to `batching="auto"`, which evolves populations against a mini-batch
+while still ranking the hall of fame on the full dataset. With `batch_size=None` the
+batch is chosen from the row count: the full dataset at N <= 1000, 128 below 5000,
+256 below 50,000, and 512 above. This problem has 20,000 rows, so evolution sees 256
+rows per comparison and the hall of fame is still scored on all 20,000.
+
+Setting `batching=False` runs every population comparison against all 20,000 rows.
+That is the same search on the same data, and it is far slower.
+"""
+
+import numpy as np
+
+from pysr import PySRRegressor
+
+rng = np.random.default_rng(0)
+X = rng.uniform(-3, 3, (20000, 5))
+y = 2.5382 * np.cos(X[:, 3]) + X[:, 0] ** 2 - 0.5
+
+VARIABLE_NAMES = ["x0", "x1", "x2", "x3", "x4"]
+
+MODEL_KWARGS = dict(
+ binary_operators=["+", "-", "*", "/"],
+ unary_operators=["cos", "exp"],
+ batching="auto",
+ niterations=40,
+ verbosity=0,
+)
+
+
+def check(model):
+ """The target is recovered when a front member reaches numerically zero loss."""
+ return bool(model.equations_["loss"].min() < 1e-10)
+
+
+def main():
+ model = PySRRegressor(**MODEL_KWARGS, random_state=0)
+ model.fit(X, y, variable_names=VARIABLE_NAMES)
+ print(model.equations_[["complexity", "loss", "equation"]].to_string(index=False))
+ print(f"\nbest loss: {model.equations_['loss'].min():.3e}")
+ print(f"recovered: {check(model)}")
+
+
+if __name__ == "__main__":
+ main()
diff --git a/examples/backsolve_mutation.py b/examples/backsolve_mutation.py
new file mode 100644
index 000000000..2ea847568
--- /dev/null
+++ b/examples/backsolve_mutation.py
@@ -0,0 +1,123 @@
+"""Recover 2.5*cos(3x) + 0.5x^2 - 1 with the backsolve mutation switched on.
+
+Backsolve inverts the operators sitting above a chosen subtree, which turns
+"improve this whole expression" into "solve for the piece that belongs in this
+hole", then fits that piece as a sparse weighted sum of parts already present in
+the population. It is off by default (`weight_backsolve` maps to weight 0.0);
+this target is out of reach for the ordinary mutations alone.
+"""
+
+import numpy as np
+
+from pysr import PySRRegressor
+
+x = np.linspace(-3.0, 3.0, 200)
+X = x.reshape(-1, 1)
+y = 2.5 * np.cos(3.0 * x) + 0.5 * x * x - 1.0
+VARIABLE_NAMES = ["x"]
+
+MODEL_KWARGS = dict(
+ binary_operators=["+", "-", "*"],
+ unary_operators=["cos", "exp"],
+ weight_backsolve=1.0,
+ maxsize=30,
+ precision=64,
+ niterations=10,
+ verbosity=0,
+)
+
+# One backsolve event, applied outside the search loop so its two losses can be
+# read directly. `mutate!` returns a member only when its monotone gate has
+# already confirmed the child beats the parent on the data.
+_BACKSOLVE_EVENT = """
+function backsolve_event(dataset, population, options, mutation)
+ events = NamedTuple[]
+ for member in population.members
+ tree = copy(member.tree)
+ result = SymbolicRegression.mutate!(
+ tree, member, mutation, options;
+ trace=nothing, dataset, curmaxsize=options.maxsize,
+ nfeatures=size(dataset.X, 1), population_for_backsolve=population,
+ )
+ result.member === nothing && continue
+ push!(events, (
+ before=Float64(member.loss),
+ after=Float64(result.member.loss),
+ parent=SymbolicRegression.string_tree(member.tree, options),
+ child=SymbolicRegression.string_tree(result.member.tree, options),
+ ))
+ end
+ return events
+end
+"""
+
+
+def check(model):
+ """The target recovered to double-precision round-off."""
+ prediction = np.asarray(model.predict(X), dtype=float)
+ return bool(np.max(np.abs(prediction - y)) < 1e-10)
+
+
+def backsolve_events(model, seed=0):
+ """Apply one backsolve mutation to every member of the search's populations.
+
+ Returns the events the mutation's own acceptance gate let through, each with
+ the parent's loss before and the child's loss after.
+ """
+ from pysr.julia_import import SymbolicRegression as SR
+ from pysr.julia_import import jl
+
+ jl.seval("using Random: seed!")
+ jl.seval(f"seed!({seed})")
+ run_event = jl.seval(_BACKSOLVE_EVENT)
+
+ dtype = np.float64 if MODEL_KWARGS["precision"] == 64 else np.float32
+ dataset = SR.Dataset(
+ jl.Matrix(np.asfortranarray(X.T.astype(dtype))), jl.Vector(y.astype(dtype))
+ )
+ mutation = SR.BacksolveMutation()
+ options = model.julia_options_
+ populations = model.julia_state_[0][0]
+
+ events, swept = [], 0
+ for i in range(int(jl.length(populations))):
+ population = populations[i]
+ swept += int(population.n)
+ found = run_event(dataset, population, options, mutation)
+ for k in range(int(jl.length(found))):
+ e = found[k]
+ events.append(
+ dict(
+ before=float(e.before),
+ after=float(e.after),
+ parent=str(e.parent),
+ child=str(e.child),
+ )
+ )
+ return events, swept
+
+
+def main():
+ model = PySRRegressor(**MODEL_KWARGS, random_state=0)
+ model.fit(X, y, variable_names=VARIABLE_NAMES)
+ print(model.equations_[["complexity", "loss", "equation"]].to_string(index=False))
+
+ prediction = np.asarray(model.predict(X), dtype=float)
+ print(f"\nmax abs error: {np.max(np.abs(prediction - y)):.3g}")
+ print(f"recovered: {check(model)}")
+
+ events, swept = backsolve_events(model)
+ print(f"\nbacksolve events accepted: {len(events)} of {swept} members swept")
+ if events:
+ ratios = sorted(e["after"] / e["before"] for e in events if e["before"] > 0)
+ print(f"median child/parent loss ratio: {ratios[len(ratios) // 2]:.3g}")
+ best = min(
+ events, key=lambda e: e["after"] / e["before"] if e["before"] > 0 else 1.0
+ )
+ print(f"largest drop: loss {best['before']:.6g} -> {best['after']:.6g}")
+ print(f" parent: {best['parent']}")
+ print(f" child: {best['child']}")
+
+
+if __name__ == "__main__":
+ main()
diff --git a/examples/cartpole_objective.py b/examples/cartpole_objective.py
new file mode 100644
index 000000000..a1842ebf8
--- /dev/null
+++ b/examples/cartpole_objective.py
@@ -0,0 +1,223 @@
+"""Search for a cart-pole swing-up policy through a custom expression-level objective.
+
+There is no target column. The loss takes a whole candidate expression, drives a cart-pole
+plant with it for ten simulated seconds from each of sixteen starts, and returns minus the
+mean per-step reward. The reward pays for an upright pole and charges for cart offset,
+motor effort and jerk. A rollout that leaves the numerical-divergence box scores a
+sentinel.
+
+Most of the starts hang the pole straight down, so a policy has to pump energy in, catch
+the pole at the top and hold it there, all from one closed-form expression of the five
+observations. The expression returns a force in units of the ten-newton actuator cap.
+
+Reward-shaped losses go negative, which the default logarithmic complexity scaling cannot
+represent, so the search uses linear loss scaling.
+
+Success is the champion earning positive mean reward on all sixty-four held-out starts,
+which are drawn wider than the training starts in every coordinate.
+"""
+
+import numpy as np
+import sympy
+
+from pysr import PySRRegressor
+
+M_C, M_P, LENGTH, GRAVITY = 1.0, 0.1, 0.5, 9.8
+FORCE_CAP, DT = 10.0, 0.02
+RAIL, V_SCALE, OMEGA_SCALE = 2.4, 4.0, 8.0
+HORIZON = 500 # 10 s of control at dt = 0.02
+
+
+def _starts(seed, rows, groups):
+ """Starts drawn one coordinate at a time, in the order the film registered them."""
+ rng = np.random.Generator(np.random.PCG64DXSM(seed))
+ out = [(x, 0.0, np.pi, 0.0) for x in rows]
+ for count, x_limit, centre, theta_limit, omega_limit in groups:
+ for _ in range(count):
+ out.append(
+ (
+ rng.uniform(-x_limit, x_limit),
+ rng.uniform(-x_limit, x_limit),
+ centre + rng.uniform(-theta_limit, theta_limit),
+ rng.uniform(-omega_limit, omega_limit),
+ )
+ )
+ return np.asarray(out, dtype=float)
+
+
+TRAIN_STARTS = _starts(
+ 20260819,
+ (0.0, 0.25),
+ (
+ (6, 0.35, np.pi, 0.35, 0.5),
+ (2, 0.35, np.pi / 2, 0.20, 0.5),
+ (2, 0.35, -np.pi / 2, 0.20, 0.5),
+ (4, 0.20, 0.0, 0.18, 0.5),
+ ),
+)
+HELD_OUT_STARTS = _starts(
+ 20260820,
+ (-0.40, -0.15, 0.15, 0.40),
+ (
+ (28, 0.55, np.pi, 0.70, 1.0),
+ (8, 0.55, np.pi / 2, 0.35, 1.0),
+ (8, 0.55, -np.pi / 2, 0.35, 1.0),
+ (16, 0.30, 0.0, 0.25, 1.0),
+ ),
+)
+
+
+def step(state, force):
+ """One semi-implicit Euler action interval for a batch of cart-pole states."""
+ x, v, theta, omega = state.T
+ total = M_C + M_P
+ sin, cos = np.sin(theta), np.cos(theta)
+ q = (force + M_P * LENGTH * omega**2 * sin) / total
+ theta_dd = (GRAVITY * sin - cos * q) / (LENGTH * (4 / 3 - M_P * cos**2 / total))
+ x_dd = q - M_P * LENGTH * theta_dd * cos / total
+ v, omega = v + DT * x_dd, omega + DT * theta_dd
+ return np.stack([x + DT * v, v, theta + DT * omega, omega], axis=-1)
+
+
+def observe(state):
+ """The five policy inputs: scaled cart offset and speed, the pole as a unit vector,
+ and scaled pole rate. They encode the four physical coordinates without loss."""
+ x, v, theta, omega = state.T
+ return np.stack(
+ [x / RAIL, v / V_SCALE, np.sin(theta), np.cos(theta), omega / OMEGA_SCALE],
+ axis=-1,
+ )
+
+
+X = observe(TRAIN_STARTS)
+y = np.zeros(len(X))
+VARIABLE_NAMES = ["x_n", "v_n", "s", "c", "omega_n"]
+
+CARTPOLE_REWARD = """
+function cartpole_reward(ex, dataset::Dataset{T,L}, options)::L where {T,L}
+ n = size(dataset.X, 2)
+ x = %(rail)r .* Float64.(dataset.X[1, :])
+ v = %(v_scale)r .* Float64.(dataset.X[2, :])
+ theta = atan.(Float64.(dataset.X[3, :]), Float64.(dataset.X[4, :]))
+ omega = %(omega_scale)r .* Float64.(dataset.X[5, :])
+
+ obs = Matrix{Float64}(undef, 5, n)
+ previous = zeros(Float64, n)
+ total = %(m_c)r + %(m_p)r
+ earned = 0.0
+ for _ in 1:%(horizon)d
+ sin_t, cos_t = sin.(theta), cos.(theta)
+ obs[1, :] .= x ./ %(rail)r
+ obs[2, :] .= v ./ %(v_scale)r
+ obs[3, :] .= sin_t
+ obs[4, :] .= cos_t
+ obs[5, :] .= omega ./ %(omega_scale)r
+ raw, ok = eval_tree_array(ex, obs, options)
+ # A batched evaluation is only usable if it returned one force per start.
+ (ok && length(raw) == n && all(isfinite, raw)) || return L(Inf)
+
+ force = %(cap)r .* clamp.(raw, -1.0, 1.0)
+ earned += sum(
+ 2.0 .* cos_t .- 0.05 .* x .^ 2 .-
+ 0.01 .* (force ./ %(cap)r) .^ 2 .-
+ 0.005 .* ((force .- previous) ./ %(cap)r) .^ 2
+ )
+ previous = force
+
+ q = (force .+ %(m_p)r * %(length)r .* omega .^ 2 .* sin_t) ./ total
+ theta_dd = (%(gravity)r .* sin_t .- cos_t .* q) ./
+ (%(length)r .* (4 / 3 .- %(m_p)r .* cos_t .^ 2 ./ total))
+ x_dd = q .- %(m_p)r * %(length)r .* theta_dd .* cos_t ./ total
+ v .+= %(dt)r .* x_dd
+ omega .+= %(dt)r .* theta_dd
+ x .+= %(dt)r .* v
+ theta .+= %(dt)r .* omega
+
+ all(isfinite, x) && all(isfinite, v) && all(isfinite, omega) &&
+ maximum(abs, x) <= 20.0 && maximum(abs, v) <= 100.0 &&
+ maximum(abs, omega) <= 100.0 || return L(1.0e12)
+ end
+ return L(-earned / (%(horizon)d * n))
+end
+""" % dict(
+ m_c=M_C,
+ m_p=M_P,
+ length=LENGTH,
+ gravity=GRAVITY,
+ cap=FORCE_CAP,
+ dt=DT,
+ rail=RAIL,
+ v_scale=V_SCALE,
+ omega_scale=OMEGA_SCALE,
+ horizon=HORIZON,
+)
+
+MODEL_KWARGS = dict(
+ operators={
+ 1: ["square", "abs", "tanh"],
+ 2: ["+", "-", "*", "/", "max", "min"],
+ 3: ["ifelse(t, a, b) = t > 0 ? a : b"],
+ },
+ extra_sympy_mappings={
+ "ifelse": lambda t, a, b: sympy.Piecewise((a, t > 0), (b, True))
+ },
+ loss_function_expression=CARTPOLE_REWARD,
+ loss_scale="linear", # rewards make the loss negative, which log scaling forbids
+ maxsize=28,
+ maxdepth=10,
+ parsimony=0.001,
+ niterations=120,
+ populations=32,
+ population_size=64,
+ verbosity=0,
+)
+
+
+def mean_rewards(model, starts, index=None):
+ """Per-start mean reward over the full horizon, with the training divergence rule."""
+ state = np.array(starts, dtype=float)
+ previous = np.zeros(len(state))
+ earned = np.zeros(len(state))
+ alive = np.ones(len(state), dtype=bool)
+ for _ in range(HORIZON):
+ raw = np.asarray(model.predict(observe(state), index=index), dtype=float)
+ force = FORCE_CAP * np.clip(np.nan_to_num(raw, nan=0.0), -1.0, 1.0)
+ earned += np.where(
+ alive,
+ 2.0 * np.cos(state[:, 2])
+ - 0.05 * state[:, 0] ** 2
+ - 0.01 * (force / FORCE_CAP) ** 2
+ - 0.005 * ((force - previous) / FORCE_CAP) ** 2,
+ 0.0,
+ )
+ previous = force
+ state = step(state, force)
+ alive &= np.isfinite(state).all(axis=-1)
+ alive &= np.abs(state[:, 0]) <= 20.0
+ alive &= np.abs(state[:, 1]) <= 100.0
+ alive &= np.abs(state[:, 3]) <= 100.0
+ state = np.nan_to_num(state, nan=0.0, posinf=0.0, neginf=0.0)
+ return np.where(alive, earned / HORIZON, -1.0e12)
+
+
+def check(model):
+ """The champion earns positive mean reward on every one of the 64 held-out starts."""
+ index = int(model.equations_["loss"].idxmin())
+ return bool((mean_rewards(model, HELD_OUT_STARTS, index=index) > 0.0).all())
+
+
+def main():
+ model = PySRRegressor(**MODEL_KWARGS, random_state=0)
+ model.fit(X, y, variable_names=VARIABLE_NAMES)
+ print(model.equations_[["complexity", "loss", "equation"]].to_string(index=False))
+
+ index = int(model.equations_["loss"].idxmin())
+ row = model.equations_.loc[index]
+ print(f"\nchampion (complexity {row.complexity}, loss {row.loss:+.6f}):")
+ print(f" {row.equation}")
+ returns = mean_rewards(model, HELD_OUT_STARTS, index=index)
+ print(f"held-out positive: {int((returns > 0).sum())}/64 starts")
+
+
+if __name__ == "__main__":
+ main()
diff --git a/examples/conftest.py b/examples/conftest.py
new file mode 100644
index 000000000..06ab2e95d
--- /dev/null
+++ b/examples/conftest.py
@@ -0,0 +1,30 @@
+"""Marker registration for the documentation-example suite.
+
+The slow examples are real searches measured in minutes, so they are opt-in. A bare
+`pytest` run exercises every example that finishes in about a minute.
+"""
+
+
+def pytest_addoption(parser):
+ parser.addoption(
+ "--runslow",
+ action="store_true",
+ default=False,
+ help="also run the examples whose searches take minutes",
+ )
+
+
+def pytest_configure(config):
+ config.addinivalue_line("markers", "slow: a search measured in minutes")
+
+
+def pytest_collection_modifyitems(config, items):
+ if config.getoption("--runslow"):
+ return
+
+ import pytest
+
+ skip = pytest.mark.skip(reason="needs --runslow")
+ for item in items:
+ if "slow" in item.keywords:
+ item.add_marker(skip)
diff --git a/examples/game_of_life.py b/examples/game_of_life.py
new file mode 100644
index 000000000..48cc74fd5
--- /dev/null
+++ b/examples/game_of_life.py
@@ -0,0 +1,58 @@
+"""Rediscover Conway's Game of Life (B3/S23) from its transition table.
+
+The search sees every combination of a cell state and a neighbour count, and nothing
+else. Its vocabulary is and / or / not plus an integer equality test; the neighbour
+thresholds are integer constants it has to find.
+"""
+
+import numpy as np
+
+from pysr import PySRRegressor, TypeSpec
+
+TABLE = [(a, n, int(n == 3 or (a == 1 and n == 2))) for a in (0, 1) for n in range(9)]
+REPLICAS = 8
+X = np.array([[a, n] for a, n, _ in TABLE] * REPLICAS, dtype=object)
+y = np.array([t for _, _, t in TABLE] * REPLICAS, dtype=object)
+VARIABLE_NAMES = ["alive", "n"]
+
+SPEC = TypeSpec(
+ "Cell",
+ fields={"v": "Int"},
+ sample="rng -> Cell(rand(rng, 0:8))",
+ mutate="(rng, value, temperature) -> Cell(mod(value.v + rand(rng, (-1, 1)), 9))",
+ string="value -> string(value.v)",
+)
+
+MODEL_KWARGS = dict(
+ type_spec=SPEC,
+ operators={
+ 1: ["not(a::Cell) = Cell(a.v == 0 ? 1 : 0)"],
+ 2: [
+ "and(a::Cell, b::Cell) = Cell((a.v != 0 && b.v != 0) ? 1 : 0)",
+ "or(a::Cell, b::Cell) = Cell((a.v != 0 || b.v != 0) ? 1 : 0)",
+ "eq(a::Cell, b::Cell) = Cell(a.v == b.v ? 1 : 0)",
+ ],
+ },
+ elementwise_loss="cell_loss(prediction::Cell, target::Cell)::Float64 = prediction.v == target.v ? 0.0 : 1.0",
+ niterations=160,
+ verbosity=0,
+)
+
+
+def check(model):
+ """Every row of the transition table reproduced exactly."""
+ prediction = np.array([int(v) for v in model.predict(X)])
+ return bool((prediction == np.array([int(v) for v in y])).all())
+
+
+def main():
+ model = PySRRegressor(**MODEL_KWARGS, random_state=0)
+ model.fit(X, y, variable_names=VARIABLE_NAMES)
+ print(model.equations_[["complexity", "loss", "equation"]].to_string(index=False))
+ prediction = np.array([int(v) for v in model.predict(X)])
+ target = np.array([int(v) for v in y])
+ print(f"\nexact rows: {int((prediction == target).sum())}/{len(target)}")
+
+
+if __name__ == "__main__":
+ main()
diff --git a/examples/hopper_controller.py b/examples/hopper_controller.py
new file mode 100644
index 000000000..c4da64105
--- /dev/null
+++ b/examples/hopper_controller.py
@@ -0,0 +1,119 @@
+"""Replay the hopping controller PySR evolved against MuJoCo rollouts.
+
+Every candidate expression was scored by simulating the Gymnasium Hopper-v5 and
+accumulating its reward, so the search optimised locomotion directly. There was
+no expert policy and no imitation data. The search is a cluster job; this script
+replays what it found, against the twenty held-out start states that were frozen
+before the search and never scored during it.
+
+Requires `gymnasium[mujoco]`.
+"""
+
+import gymnasium as gym
+import numpy as np
+
+HELDOUT_SEEDS = range(20261001, 20261021)
+RESET_NOISE = 0.01
+
+# Eight seconds at the Hopper's 0.008 s control step. Passing means staying
+# upright for all of it, travelling five metres, and completing four
+# contact/liftoff cycles.
+STEPS = 1000
+DISTANCE = 5.0
+CYCLES = 4
+
+# The three vector constants the search chose, one component per motor.
+C0 = np.array([-0.24771142177432096, -4.22684590927177286, 8.8397066748254275])
+C1 = np.array([0.06452662644513463, 0.6180986305329657, 0.3968791265257595])
+C2 = np.array([0.49735344234321904, -1.1920811500492086, 1.1808251376288732])
+
+
+def action(qpos, qvel):
+ """The evolved controller: hopper state in, three motor commands out."""
+ angles = qpos[3:6]
+ rates = np.clip(qvel[3:6], -10.0, 10.0) / 10.0
+ pitch = qpos[2]
+ forward, vertical, spin = qvel[0] / 5.0, qvel[1] / 5.0, qvel[2] / 5.0
+ rolled_angles = angles[[1, 2, 0]]
+ rolled_rates = rates[[1, 2, 0]]
+ d = angles - 2 * rolled_angles + rolled_rates - 2 * (forward + spin) + C0
+ left = -angles - (pitch + spin * (rolled_rates - spin)) * d - C1
+ right = -rolled_angles - vertical + spin + C2
+ return np.tanh(rolled_angles - spin + left * right)
+
+
+def contact_geoms(env):
+ model = env.unwrapped.model
+ named = [(g, (model.geom(g).name or "").lower()) for g in range(model.ngeom)]
+ return (
+ {g for g, name in named if "foot" in name},
+ {g for g, name in named if "floor" in name},
+ )
+
+
+def touching(data, feet, floor):
+ pairs = ((int(c.geom1), int(c.geom2)) for c in data.contact[: data.ncon])
+ return any(a in feet and b in floor or b in feet and a in floor for a, b in pairs)
+
+
+def cycles(contact):
+ """Completed liftoff-to-landing cycles, which is what makes a gait a gait."""
+ count = 0
+ landed = airborne = False
+ for down in contact:
+ if down:
+ count += airborne
+ landed, airborne = True, False
+ elif landed:
+ airborne = True
+ return count
+
+
+def trial(env, feet, floor, seed):
+ env.reset(seed=seed)
+ data = env.unwrapped.data
+ start = float(data.qpos[0])
+ contact = []
+ upright = True
+ for _ in range(STEPS):
+ _, _, terminated, truncated, _ = env.step(action(data.qpos, data.qvel))
+ contact.append(touching(data, feet, floor))
+ upright = not terminated
+ if terminated or truncated:
+ break
+ hops = cycles(contact)
+ distance = float(data.qpos[0]) - start
+ return {
+ "seed": seed,
+ "steps": len(contact),
+ "distance": distance,
+ "cycles": hops,
+ "success": (
+ len(contact) == STEPS
+ and upright
+ and distance >= DISTANCE
+ and hops >= CYCLES
+ ),
+ }
+
+
+def replay():
+ env = gym.make("Hopper-v5", reset_noise_scale=RESET_NOISE)
+ feet, floor = contact_geoms(env)
+ try:
+ return [trial(env, feet, floor, seed) for seed in HELDOUT_SEEDS]
+ finally:
+ env.close()
+
+
+def main():
+ results = replay()
+ for r in results:
+ outcome = "passed" if r["success"] else "fell"
+ line = f"seed {r['seed']} {outcome:<6} {r['distance']:6.2f} m"
+ print(f"{line} {r['cycles']:2d} hops {r['steps']:4d} steps")
+ print(f"\n{sum(r['success'] for r in results)} of {len(results)} held-out starts")
+
+
+if __name__ == "__main__":
+ main()
diff --git a/examples/magnetic_field.py b/examples/magnetic_field.py
new file mode 100644
index 000000000..21ca31c16
--- /dev/null
+++ b/examples/magnetic_field.py
@@ -0,0 +1,146 @@
+"""Recover a magnetic field from force measurements alone.
+
+A charged particle obeys F = -eta(T) v + v x B(t), and 1000 experiments record only
+the total force and the inputs (t, v, T). A TemplateExpressionSpec supplies that
+structure and a Force type spec carries the three-component data, so the search has
+only to fill in B_x(t), B_y(t), B_z(t) and the drag scale eta(T).
+"""
+
+import numpy as np
+
+from pysr import PySRRegressor, TemplateExpressionSpec, TypeSpec
+
+OMEGA = 2 * np.pi
+
+
+def experiments(n, seed):
+ rng = np.random.default_rng(seed)
+ T = 298.15 + 0.5 * rng.random(n)
+ t = 10.0 * rng.random(n)
+ v = 2.0 * rng.random((n, 3)) - 1.0
+ B = np.stack([np.sin(OMEGA * t), np.cos(OMEGA * t), np.exp(-t / 10.0)], axis=1)
+ F = -1e-5 * np.sqrt(T)[:, None] * v + np.cross(v, B)
+ return np.column_stack([t, v[:, 0], v[:, 1], v[:, 2], T]), F
+
+
+def as_forces(rows):
+ """An (n, 3) float array as a column of n Force values."""
+ out = np.empty(len(rows), dtype=object)
+ out[:] = [tuple(map(float, row)) for row in rows]
+ return out
+
+
+def on_diagonal(columns):
+ """Every value in the problem is a Force, so scalar inputs ride the diagonal."""
+ out = np.empty(columns.shape, dtype=object)
+ for j in range(columns.shape[1]):
+ out[:, j] = as_forces(np.repeat(columns[:, j, None], 3, axis=1))
+ return out
+
+
+INPUTS, FORCES = experiments(1000, seed=0)
+X = on_diagonal(INPUTS)
+y = as_forces(FORCES)
+VARIABLE_NAMES = ["t", "v_x", "v_y", "v_z", "T"]
+
+HELD_OUT_INPUTS, HELD_OUT_FORCES = experiments(500, seed=12345)
+HELD_OUT_X = on_diagonal(HELD_OUT_INPUTS)
+
+FORCE = TypeSpec(
+ "Force",
+ fields={"x": "Float64", "y": "Float64", "z": "Float64"},
+ sample="rng -> Force(randn(rng, 3)...)",
+ definitions="""
+ Base.sin(a::Force)::Force = Force(sin(a.x), sin(a.y), sin(a.z))
+ Base.cos(a::Force)::Force = Force(cos(a.x), cos(a.y), cos(a.z))
+ function Base.sqrt(a::Force)::Force
+ f(v) = v < 0 ? NaN : sqrt(v)
+ return Force(f(a.x), f(a.y), f(a.z))
+ end
+ Base.exp(a::Force)::Force = Force(exp(a.x), exp(a.y), exp(a.z))
+ Base.:+(a::Force, b::Force)::Force = Force(a.x + b.x, a.y + b.y, a.z + b.z)
+ Base.:-(a::Force, b::Force)::Force = Force(a.x - b.x, a.y - b.y, a.z - b.z)
+ Base.:*(a::Force, b::Force)::Force = Force(a.x * b.x, a.y * b.y, a.z * b.z)
+ Base.:/(a::Force, b::Force)::Force = Force(a.x / b.x, a.y / b.y, a.z / b.z)
+ """,
+ scalar_constants="value -> [value.x, value.y, value.z]",
+ with_scalar_constants="(value, c) -> Force(c[1], c[2], c[3])",
+ string="value -> sprint(show, (value.x, value.y, value.z); context = :compact => true)",
+ init_invalid="() -> Force(NaN, NaN, NaN)",
+)
+
+
+OPERATORS = {
+ 1: [
+ "Base.sin",
+ "Base.cos",
+ "Base.sqrt",
+ "Base.exp",
+ ],
+ 2: [
+ "Base.:+",
+ "Base.:-",
+ "Base.:*",
+ "Base.:/",
+ ],
+}
+
+FORCE_LOSS = "force_loss(a::Force, b::Force)::Float64 = (a.x - b.x)^2 + (a.y - b.y)^2 + (a.z - b.z)^2"
+
+PHYSICS = r"""
+begin
+ _B_x = B_x(t)
+ _B_y = B_y(t)
+ _B_z = B_z(t)
+ _F_d_scale = F_d_scale(T)
+ if !(_B_x.valid && _B_y.valid && _B_z.valid && _F_d_scale.valid)
+ return ValidVector(_B_x.x, false)
+ end
+ F = map(_B_x.x, _B_y.x, _B_z.x, _F_d_scale.x,
+ v_x.x, v_y.x, v_z.x) do bx, by, bz, fd, ux, uy, uz
+ b1, b2, b3 = bx.x, by.x, bz.x
+ u1, u2, u3 = ux.x, uy.x, uz.x
+ s = fd.x
+ Force(u2 * b3 - u3 * b2 + s * u1,
+ u3 * b1 - u1 * b3 + s * u2,
+ u1 * b2 - u2 * b1 + s * u3)
+ end
+ ValidVector(F, true)
+end
+"""
+
+STRUCTURE = TemplateExpressionSpec(
+ combine=PHYSICS,
+ expressions=["B_x", "B_y", "B_z", "F_d_scale"],
+ variable_names=VARIABLE_NAMES,
+)
+
+MODEL_KWARGS = dict(
+ type_spec=FORCE,
+ expression_spec=STRUCTURE,
+ operators=OPERATORS,
+ elementwise_loss=FORCE_LOSS,
+ niterations=100,
+ verbosity=0,
+)
+
+
+def held_out_error(model):
+ predicted = np.array([list(f) for f in model.predict(HELD_OUT_X)], dtype=float)
+ return float(np.abs(predicted - HELD_OUT_FORCES).max())
+
+
+def check(model):
+ """Forces on unseen experiments reproduced to machine precision."""
+ return held_out_error(model) < 1e-6
+
+
+def main():
+ model = PySRRegressor(**MODEL_KWARGS, random_state=0)
+ model.fit(X, y, variable_names=VARIABLE_NAMES)
+ print(model.equations_[["complexity", "loss", "equation"]].to_string(index=False))
+ print(f"\nheld-out max abs error: {held_out_error(model):.3g}")
+
+
+if __name__ == "__main__":
+ main()
diff --git a/examples/morton_interleave.py b/examples/morton_interleave.py
new file mode 100644
index 000000000..87f658b44
--- /dev/null
+++ b/examples/morton_interleave.py
@@ -0,0 +1,114 @@
+"""Rediscover a bit-exact Morton (Z-order) interleave of two 4-bit integers.
+
+The search works on a machine word: a Word value backed by UInt32 whose vocabulary is
+the integer instruction set (xor, and, or, not, shift left, shift right), with shift
+distances as ordinary Word constants it has to find. It scores all 256 four-bit input
+pairs and has to recover the map that interleaves their bits.
+
+The score only ever sees those 256 pairs, so what comes back is an expression that
+agrees with the interleave everywhere it was asked. That is a weaker claim than the
+general Z-order map, and the difference is measurable: the recovered forms are exact on
+every 4-bit pair and correct on only 256 of the 1024 5-bit pairs, which are exactly the
+4-bit ones. The run prints both numbers.
+
+Two things about the shape of the answer are worth knowing before reading it. Written in
+this vocabulary as a tree, with no shared subexpressions, the textbook interleave kernel
+(two shift-mask spread stages per input, then one shift and or) is 41 nodes, so a
+recovered form near 40 is the same size as the hand-written one rather than bloat, and
+maxsize has to leave room for it. And the masks print as arbitrary hex because only
+their low bits can affect an 8-bit result: in 0x00080022 the search pinned bits 1 and 5
+and bit 19 is free, since no scored pair can tell what it holds.
+"""
+
+import numpy as np
+
+from pysr import PySRRegressor, TypeSpec
+
+WIDTH = 4
+WIDER = 5
+
+
+def morton(x, y, width=WIDTH):
+ z = 0
+ for i in range(width):
+ z |= ((x >> i) & 1) << (2 * i)
+ z |= ((y >> i) & 1) << (2 * i + 1)
+ return z
+
+
+PAIRS = [(a, b) for b in range(1 << WIDTH) for a in range(1 << WIDTH)]
+X = np.array(PAIRS, dtype=object)
+y = np.array([morton(a, b) for a, b in PAIRS], dtype=object)
+VARIABLE_NAMES = ["x", "y"]
+
+SPEC = TypeSpec(
+ "Word",
+ fields={"bits": "UInt32"},
+ sample="rng -> Word(rand(rng, Bool) ? rand(rng, UInt32(0):UInt32(31)) : rand(rng, UInt32))",
+ mutate="(rng, value, temperature) -> Word(rand(rng, Bool) ? xor(value.bits, UInt32(1) << rand(rng, 0:31)) : value.bits + rand(rng, (UInt32(1), typemax(UInt32))))",
+ string='value -> value.bits < UInt32(32) ? string(value.bits) : "0x" * string(value.bits, base = 16, pad = 8)',
+)
+
+MODEL_KWARGS = dict(
+ type_spec=SPEC,
+ operators={
+ 1: ["bnot(a::Word) = Word(~a.bits)"],
+ 2: [
+ "bxor(a::Word, b::Word) = Word(xor(a.bits, b.bits))",
+ "band(a::Word, b::Word) = Word(a.bits & b.bits)",
+ "bor(a::Word, b::Word) = Word(a.bits | b.bits)",
+ "shl(a::Word, b::Word) = Word(a.bits << b.bits)",
+ "shr(a::Word, b::Word) = Word(a.bits >> b.bits)",
+ ],
+ },
+ elementwise_loss="bit_loss(prediction::Word, target::Word)::Float64 = count_ones(xor(prediction.bits, target.bits)) / 32",
+ maxsize=45,
+ niterations=800,
+ verbosity=0,
+)
+
+
+def exact_front_rows(model):
+ """Indices of front members that reproduce every one of the 256 interleaves."""
+ target = np.array([int(v) for v in y])
+ rows = []
+ for i in range(len(model.equations_)):
+ prediction = np.array([int(v) for v in model.predict(X, index=i)])
+ if (prediction == target).all():
+ rows.append(i)
+ return rows
+
+
+def check(model):
+ """Success is an exact interleaver anywhere on the Pareto front."""
+ return bool(exact_front_rows(model))
+
+
+def wider_agreement(model, index):
+ """How many 5-bit pairs the recovered expression interleaves correctly."""
+ pairs = [(a, b) for b in range(1 << WIDER) for a in range(1 << WIDER)]
+ inputs = np.array(pairs, dtype=object)
+ expected = np.array([morton(a, b, WIDER) for a, b in pairs])
+ got = np.array([int(v) for v in model.predict(inputs, index=index)])
+ return int((got == expected).sum()), len(pairs)
+
+
+def main():
+ model = PySRRegressor(**MODEL_KWARGS, random_state=0)
+ model.fit(X, y, variable_names=VARIABLE_NAMES)
+ print(model.equations_[["complexity", "loss", "equation"]].to_string(index=False))
+ rows = exact_front_rows(model)
+ if rows:
+ best = model.equations_.iloc[rows[0]]
+ print(f"\nexact interleaver at complexity {int(best.complexity)}:")
+ print(best.equation)
+ agree, total = wider_agreement(model, rows[0])
+ print(
+ f"\non {WIDER}-bit pairs it interleaves {agree}/{total} correctly, so it is a {WIDTH}-bit interleaver and not the Z-order map"
+ )
+ else:
+ print("\nno exact interleaver on the front")
+
+
+if __name__ == "__main__":
+ main()
diff --git a/examples/mutations_and_plugins.py b/examples/mutations_and_plugins.py
new file mode 100644
index 000000000..5fb5008e6
--- /dev/null
+++ b/examples/mutations_and_plugins.py
@@ -0,0 +1,64 @@
+"""Weight a mutation and add a plugin from PySRRegressor.
+
+Mutations and plugins are ordinary configuration: `mutations` maps a mutation
+configuration to a weight, `plugins` lists plugin configurations, and entries of
+either override the shipped defaults by type and leave the rest alone.
+`BacksolveMutation` ships at weight 0.0, so giving it a weight is what turns it
+on for this search for 2.5*cos(3x) + 0.5*x^2 - 1.
+"""
+
+import numpy as np
+
+from pysr import AdaptiveMutationWeightsPlugin, BacksolveMutation, PySRRegressor
+from pysr.julia_import import jl
+
+x = np.linspace(-3.0, 3.0, 200)
+X = x.reshape(-1, 1)
+y = 2.5 * np.cos(3.0 * x) + 0.5 * x * x - 1.0
+VARIABLE_NAMES = ["x"]
+
+BASE_KWARGS = dict(
+ binary_operators=["+", "-", "*"],
+ unary_operators=["cos", "exp"],
+ niterations=20,
+ verbosity=0,
+)
+
+MODEL_KWARGS = dict(
+ **BASE_KWARGS,
+ mutations={BacksolveMutation(): 0.1},
+ plugins=[AdaptiveMutationWeightsPlugin()],
+)
+
+
+def check(model):
+ """The noiseless target reproduced to within 1e-5 everywhere."""
+ return bool(np.max(np.abs(model.predict(X) - y)) < 1e-5)
+
+
+def _backsolve_weight(model):
+ for mutation, weight in model.julia_options_.mutations:
+ if "BacksolveMutation" in jl.string(mutation):
+ return float(weight)
+ raise AssertionError("BacksolveMutation missing from the resolved mutation table")
+
+
+def _run(kwargs, label):
+ model = PySRRegressor(**kwargs, random_state=0)
+ model.fit(X, y, variable_names=VARIABLE_NAMES)
+ print(f"=== {label}")
+ print(model.equations_[["complexity", "loss", "equation"]].to_string(index=False))
+ print(f"max error: {np.max(np.abs(model.predict(X) - y)):.3g}")
+ print(f"backsolve weight: {_backsolve_weight(model)}")
+ print(f"plugins: {jl.string(model.julia_options_.plugins)}\n")
+ return model
+
+
+def main():
+ _run(BASE_KWARGS, "default mutations and plugins")
+ model = _run(MODEL_KWARGS, "backsolve weighted, adaptive weights on")
+ print(f"exact recovery: {check(model)}")
+
+
+if __name__ == "__main__":
+ main()
diff --git a/examples/prng_period.py b/examples/prng_period.py
new file mode 100644
index 000000000..a266bc7fd
--- /dev/null
+++ b/examples/prng_period.py
@@ -0,0 +1,466 @@
+"""Invent a 32-bit pseudorandom generator with no target array.
+
+`y` is a column of zeros and the objective never reads it. The search sees 32 seed
+states, and the loss runs each candidate as an iterated 32-bit map and scores the
+generator that map is: how long its orbit is, how balanced its bit columns are, how
+uncorrelated successive states are, how far one flipped input bit spreads, and whether
+any bit row of the output is a shifted copy of a bit row of the input. Every term is 0
+exactly when its desideratum is met.
+
+The period term is the interesting one, because the whole state space is 4294967295
+nonzero words and nothing may walk it. The objective probes the candidate on 0, on the
+32 basis states and on 32 fixed states; if it is GF(2)-linear on that evidence, it
+builds the minimal polynomial of the first seed's Krylov sequence, and if that
+polynomial is irreducible the orbit's length is exactly the multiplicative order of x
+modulo it, obtained by dividing the prime factors out of 2^deg - 1. A map whose period
+cannot be certified is credited only with the orbit actually walked, so a candidate is
+paid for the period it can prove.
+
+`check` re-derives the claim in Python from the recovered expression, by a different
+route than the loss uses. It reads the map's GF(2) matrix out of `predict`, confirms
+the map really is linear on random states, and computes the multiplicative order of
+that matrix by repeated squaring. An order of 4294967295 forces the period: the order
+of a matrix over GF(2) divides the unit-group order of the algebra it generates, which
+is a power of two times a product of 2^d - 1 over the degrees of the minimal
+polynomial's irreducible factors, and 65537 divides 2^d - 1 only when 32 divides d. So
+the minimal polynomial is irreducible of degree 32, the state space is one copy of
+GF(2^32) under the map, and every nonzero word lies on one orbit of length 4294967295.
+"""
+
+import numpy as np
+
+from pysr import PySRRegressor, TypeSpec
+
+WIDTH = 32
+MASK = (1 << WIDTH) - 1
+FULL_PERIOD = MASK
+ORDER_PRIMES = (3, 5, 17, 257, 65537)
+
+SEEDS = [int(v) for v in np.random.default_rng(0).integers(1, 1 << WIDTH, size=32)]
+X = np.empty((len(SEEDS), 1), dtype=object)
+for i, seed in enumerate(SEEDS):
+ X[i, 0] = seed
+y = np.zeros(len(SEEDS), dtype=object)
+VARIABLE_NAMES = ["x"]
+
+SPEC = TypeSpec(
+ "Word",
+ fields={"bits": "UInt32"},
+ sample="rng -> Word(rand(rng, Bool) ? rand(rng, UInt32(0):UInt32(31)) : rand(rng, UInt32))",
+ mutate="(rng, value, temperature) -> Word(rand(rng, Bool) ? xor(value.bits, UInt32(1) << rand(rng, 0:31)) : value.bits + rand(rng, (UInt32(1), typemax(UInt32))))",
+ string='value -> value.bits < UInt32(32) ? string(value.bits) : "0x" * string(value.bits, base = 16, pad = 8)',
+ loss_type="Float64",
+)
+
+# The certificate: linear algebra over GF(2), so a period of four billion states costs
+# a few dozen 32-bit operations instead of four billion evaluations.
+PERIOD_CERTIFICATE = r"""
+const PRNG_PROBES = UInt32[UInt32(0x9e3779b9) * UInt32(k) for k in 1:32]
+
+function _prng_matvec(cols::Vector{UInt32}, v::UInt32)
+ acc = UInt32(0)
+ vv = v
+ while vv != 0
+ acc ⊻= cols[trailing_zeros(vv) + 1]
+ vv &= vv - UInt32(1)
+ end
+ return acc
+end
+
+function _prng_mulmod(a::UInt64, b::UInt64, f::UInt64, deg::Int)
+ out = UInt64(0)
+ aa, bb = a, b
+ while bb != 0
+ if bb & 1 == 1
+ out ⊻= aa
+ end
+ bb >>= 1
+ aa <<= 1
+ if (aa >> deg) & 1 == 1
+ aa ⊻= f
+ end
+ end
+ return out
+end
+
+function _prng_powmod(base::UInt64, e::UInt64, f::UInt64, deg::Int)
+ r = UInt64(1)
+ b, ee = base, e
+ while ee != 0
+ if ee & 1 == 1
+ r = _prng_mulmod(r, b, f, deg)
+ end
+ b = _prng_mulmod(b, b, f, deg)
+ ee >>= 1
+ end
+ return r
+end
+
+_prng_deg(p::UInt64) = 63 - leading_zeros(p)
+
+function _prng_gcd(a::UInt64, b::UInt64)
+ x, y = a, b
+ while y != 0
+ dy = _prng_deg(y)
+ while x != 0 && _prng_deg(x) >= dy
+ x ⊻= y << (_prng_deg(x) - dy)
+ end
+ x, y = y, x
+ end
+ return x
+end
+
+function _prng_minpoly(cols::Vector{UInt32}, v0::UInt32)
+ leads = Int[]
+ vecs = UInt32[]
+ combs = UInt64[]
+ cur = v0
+ comb = UInt64(1)
+ for _ in 0:32
+ red, rc = cur, comb
+ for i in eachindex(leads)
+ if (red >> leads[i]) & UInt32(1) == 1
+ red ⊻= vecs[i]
+ rc ⊻= combs[i]
+ end
+ end
+ if red == 0
+ return rc
+ end
+ push!(leads, 31 - leading_zeros(red))
+ push!(vecs, red)
+ push!(combs, rc)
+ cur = _prng_matvec(cols, cur)
+ comb <<= 1
+ end
+ return UInt64(0)
+end
+
+function _prng_irreducible(f::UInt64, deg::Int)
+ deg < 1 && return false
+ _prng_powmod(UInt64(2), UInt64(1) << deg, f, deg) == UInt64(2) || return false
+ d = deg
+ p = 2
+ while p * p <= d
+ if d % p == 0
+ g = _prng_gcd(
+ _prng_powmod(UInt64(2), UInt64(1) << (deg ÷ p), f, deg) ⊻ UInt64(2), f)
+ g == UInt64(1) || return false
+ while d % p == 0
+ d ÷= p
+ end
+ end
+ p += 1
+ end
+ if d > 1
+ g = _prng_gcd(
+ _prng_powmod(UInt64(2), UInt64(1) << (deg ÷ d), f, deg) ⊻ UInt64(2), f)
+ g == UInt64(1) || return false
+ end
+ return true
+end
+
+function _prng_order(f::UInt64, deg::Int)
+ order = (UInt64(1) << deg) - UInt64(1)
+ rest = order
+ p = UInt64(2)
+ while p * p <= rest
+ if rest % p == 0
+ while order % p == 0 && _prng_powmod(UInt64(2), order ÷ p, f, deg) == UInt64(1)
+ order ÷= p
+ end
+ while rest % p == 0
+ rest ÷= p
+ end
+ end
+ p += UInt64(1)
+ end
+ if rest > 1
+ while order % rest == 0 &&
+ _prng_powmod(UInt64(2), order ÷ rest, f, deg) == UInt64(1)
+ order ÷= rest
+ end
+ end
+ return order
+end
+
+function _prng_certified(tree, options, seed::UInt32)
+ m = 33 + length(PRNG_PROBES)
+ grid = Matrix{Word}(undef, 1, m)
+ grid[1, 1] = Word(UInt32(0))
+ for b in 0:31
+ grid[1, 2 + b] = Word(UInt32(1) << b)
+ end
+ for (i, p) in enumerate(PRNG_PROBES)
+ grid[1, 33 + i] = Word(p)
+ end
+ out, ok = eval_tree_array(tree, grid, options)
+ ok || return UInt64(0)
+ out[1].bits == 0 || return UInt64(0)
+
+ cols = [out[2 + b].bits for b in 0:31]
+ for (i, p) in enumerate(PRNG_PROBES)
+ _prng_matvec(cols, p) == out[33 + i].bits || return UInt64(0)
+ end
+
+ f = _prng_minpoly(cols, seed)
+ f == 0 && return UInt64(0)
+ deg = _prng_deg(f)
+ _prng_irreducible(f, deg) || return UInt64(0)
+ return _prng_order(f, deg)
+end
+"""
+
+PRNG_LOSS = PERIOD_CERTIFICATE + r"""
+function prng_loss(tree, dataset::Dataset{T,L}, options)::L where {T,L}
+ steps = 256
+ lags = (1, 2, 3, 5)
+ n = dataset.n
+ walk = Matrix{UInt32}(undef, steps, n)
+ x = copy(dataset.X)
+ for s in 1:steps
+ out, ok = eval_tree_array(tree, x, options)
+ ok || return L(Inf)
+ length(out) == n || error("prng_loss scored " * string(length(out)) *
+ " of " * string(n) * " samples; set batching=false")
+ for j in 1:n
+ walk[s, j] = out[j].bits
+ end
+ x = reshape(out, 1, n)
+ end
+
+ seen = typemax(Int)
+ for j in 1:n
+ seen = min(seen, length(Set(@view walk[:, j])))
+ end
+ certified = _prng_certified(tree, options, dataset.X[1, 1].bits)
+ proven = certified > 0 ? Float64(certified) : Float64(seen)
+ period = 1 - log2(proven) / 32
+
+ ones = zeros(Int, 32)
+ for j in 1:n, s in 1:steps
+ w = walk[s, j]
+ for b in 1:32
+ ones[b] += (w >> (b - 1)) & UInt32(1)
+ end
+ end
+ total = steps * n
+ balance = 0.0
+ for b in 1:32
+ p = ones[b] / total
+ balance += (2p - 1)^2
+ end
+ balance /= 32
+
+ autocorr = 0.0
+ joint = zeros(Int, 32)
+ head = zeros(Int, 32)
+ tail = zeros(Int, 32)
+ for lag in lags
+ fill!(joint, 0)
+ fill!(head, 0)
+ fill!(tail, 0)
+ for j in 1:n, s in 1:(steps - lag)
+ u = walk[s, j]
+ v = walk[s + lag, j]
+ for b in 1:32
+ a = (u >> (b - 1)) & UInt32(1)
+ c = (v >> (b - 1)) & UInt32(1)
+ head[b] += a
+ tail[b] += c
+ joint[b] += a & c
+ end
+ end
+ pairs = (steps - lag) * n
+ for b in 1:32
+ ph = head[b] / pairs
+ pt = tail[b] / pairs
+ cov = joint[b] / pairs - ph * pt
+ var = ph * (1 - ph) * pt * (1 - pt)
+ autocorr += var > 0 ? cov^2 / var : 0.0
+ end
+ end
+ autocorr /= 32 * length(lags)
+
+ m = length(PRNG_PROBES)
+ grid = Matrix{Word}(undef, 1, m)
+ for i in 1:m
+ grid[1, i] = Word(PRNG_PROBES[i])
+ end
+ base, bok = eval_tree_array(tree, grid, options)
+ bok || return L(Inf)
+ moved = 0
+ deps = zeros(Int, 32)
+ for b in 0:31
+ for i in 1:m
+ grid[1, i] = Word(PRNG_PROBES[i] ⊻ (UInt32(1) << b))
+ end
+ flipped, fok = eval_tree_array(tree, grid, options)
+ fok || return L(Inf)
+ touched = UInt32(0)
+ for i in 1:m
+ d = base[i].bits ⊻ flipped[i].bits
+ moved += count_ones(d)
+ touched |= d
+ end
+ for j in 1:32
+ deps[j] += (touched >> (j - 1)) & UInt32(1)
+ end
+ end
+ diffusion = max(0.0, (4.0 - moved / (32 * m)) / 4.0)^2
+ edge = 0.0
+ for j in 1:32
+ edge += max(0.0, (2.0 - deps[j]) / 2.0)^2
+ end
+ edge /= 32
+
+ shear = 0.0
+ npairs = 0
+ sjoint = zeros(Int, 32)
+ shead = zeros(Int, 32)
+ stail = zeros(Int, 32)
+ for lag in (1, 2), d in (-3, -2, -1, 1, 2, 3)
+ blo = d >= 0 ? 0 : -d
+ bhi = d >= 0 ? 31 - d : 31
+ fill!(sjoint, 0)
+ fill!(shead, 0)
+ fill!(stail, 0)
+ for j in 1:n, s in 1:(steps - lag)
+ u = walk[s, j]
+ v = walk[s + lag, j]
+ for b in blo:bhi
+ a = (u >> b) & UInt32(1)
+ c = (v >> (b + d)) & UInt32(1)
+ shead[b + 1] += a
+ stail[b + 1] += c
+ sjoint[b + 1] += a & c
+ end
+ end
+ spairs = (steps - lag) * n
+ for b in blo:bhi
+ ph = shead[b + 1] / spairs
+ pt = stail[b + 1] / spairs
+ cov = sjoint[b + 1] / spairs - ph * pt
+ var = ph * (1 - ph) * pt * (1 - pt)
+ shear += var > 0 ? cov^2 / var : 0.0
+ npairs += 1
+ end
+ end
+ shear /= npairs
+
+ return L(period + balance + autocorr + diffusion + shear + edge)
+end
+"""
+
+MODEL_KWARGS = dict(
+ type_spec=SPEC,
+ operators={
+ 1: ["bnot(a::Word) = Word(~a.bits)"],
+ 2: [
+ "bxor(a::Word, b::Word) = Word(xor(a.bits, b.bits))",
+ "band(a::Word, b::Word) = Word(a.bits & b.bits)",
+ "bor(a::Word, b::Word) = Word(a.bits | b.bits)",
+ "shl(a::Word, b::Word) = Word(a.bits << b.bits)",
+ "shr(a::Word, b::Word) = Word(a.bits >> b.bits)",
+ ],
+ },
+ loss_function=PRNG_LOSS,
+ # The objective walks every seed and indexes the sample axis, so it has to see the
+ # whole sample set; a minibatch would score a different problem.
+ batching=False,
+ maxsize=20,
+ # Eight small populations, evolved deeply: each candidate costs a 256-step walk of
+ # every seed, so the search is worth more generations rather than more members.
+ populations=8,
+ population_size=30,
+ ncycles_per_iteration=30,
+ niterations=800,
+ verbosity=0,
+)
+
+IDENTITY = [1 << b for b in range(WIDTH)]
+LINEARITY_PROBES = [
+ int(v) for v in np.random.default_rng(1).integers(0, 1 << WIDTH, size=256)
+]
+
+
+def matvec(cols, v):
+ """Image of `v` under the GF(2) matrix whose columns are `cols`."""
+ acc = 0
+ while v:
+ acc ^= cols[(v & -v).bit_length() - 1]
+ v &= v - 1
+ return acc
+
+
+def matmul(a, b):
+ return [matvec(a, col) for col in b]
+
+
+def matpow(m, e):
+ result = IDENTITY
+ while e:
+ if e & 1:
+ result = matmul(result, m)
+ m = matmul(m, m)
+ e >>= 1
+ return result
+
+
+def linear_matrix(model, index):
+ """The map's GF(2) matrix, or None if the recovered expression is not linear."""
+ probe = np.empty((1 + WIDTH + len(LINEARITY_PROBES), 1), dtype=object)
+ probe[0, 0] = 0
+ for b in range(WIDTH):
+ probe[1 + b, 0] = 1 << b
+ for i, state in enumerate(LINEARITY_PROBES):
+ probe[1 + WIDTH + i, 0] = state
+ image = [int(v) & MASK for v in model.predict(probe, index=index)]
+ if image[0] != 0:
+ return None
+ cols = image[1 : 1 + WIDTH]
+ for state, got in zip(LINEARITY_PROBES, image[1 + WIDTH :]):
+ if matvec(cols, state) != got:
+ return None
+ return cols
+
+
+def period(model, index):
+ """4294967295 when the map is a full-period generator, 0 when it is not.
+
+ The map is linear, so its order under composition is settled by matrix powers:
+ order 4294967295 forces the minimal polynomial to be irreducible of degree 32, and
+ then every nonzero word sits on a single orbit of that length.
+ """
+ cols = linear_matrix(model, index)
+ if cols is None or matpow(cols, FULL_PERIOD) != IDENTITY:
+ return 0
+ if any(matpow(cols, FULL_PERIOD // q) == IDENTITY for q in ORDER_PRIMES):
+ return 0
+ return FULL_PERIOD
+
+
+def check(model):
+ """Some row of the front runs through all 4294967295 nonzero words."""
+ return any(period(model, i) == FULL_PERIOD for i in model.equations_.index)
+
+
+def main():
+ model = PySRRegressor(**MODEL_KWARGS, random_state=0)
+ model.fit(X, y, variable_names=VARIABLE_NAMES)
+ front = model.equations_[["complexity", "loss", "equation"]].copy()
+ front["full_period"] = [
+ period(model, i) == FULL_PERIOD for i in model.equations_.index
+ ]
+ print(front.to_string(index=False))
+ full = front[front["full_period"]]
+ if len(full):
+ row = full.iloc[0]
+ print(
+ f"\nperiod {FULL_PERIOD} at complexity {int(row.complexity)}: {row.equation}"
+ )
+ print(f"full period reached: {bool(len(full))}")
+
+
+if __name__ == "__main__":
+ main()
diff --git a/examples/pysr_demo.ipynb b/examples/pysr_demo.ipynb
index fd7c264b7..2f2fa51bb 100644
--- a/examples/pysr_demo.ipynb
+++ b/examples/pysr_demo.ipynb
@@ -138,7 +138,7 @@
"id": "N4gANfkaj8ie"
},
"source": [
- "PySR can run for arbitrarily long, and continue to find more and more accurate expressions. You can set the total number of cycles of evolution with `niterations`, although there are also a [few more ways](https://github.com/MilesCranmer/PySR/pull/134) to stop execution.\n",
+ "PySR can run for arbitrarily long, and continue to find more and more accurate expressions. You can set the total number of cycles of evolution with `niterations`, although there are also a [few more ways](https://github.com/astroautomata/PySR/pull/134) to stop execution.\n",
"\n",
"**This first execution will take a bit longer to startup, as the library is JIT-compiled. The next execution will be much faster.**"
]
@@ -321,7 +321,7 @@
"id": "qvgVbOoSFtQY"
},
"source": [
- "A full list of operators is given here: https://ai.damtp.cam.ac.uk/pysr/operators,\n",
+ "A full list of operators is given here: https://pysr.ai/operators,\n",
"but we can also use any binary or unary operator in `julia`, or define our own as arbitrary functions.\n",
"\n",
"Say that we want a command to do quartic powers:\n",
@@ -801,8 +801,8 @@
"cell_type": "markdown",
"metadata": {},
"source": [
- "Note that PySR should automatically load the `juliacall.ipython` extension for you,\n",
- "which means that you can also execute Julia code in the notebook using the `%%julia` magic:"
+ "You can also execute Julia code in the notebook using the `%%julia` magic.\n",
+ "Load the `juliacall` IPython extension to enable it:"
]
},
{
@@ -811,6 +811,7 @@
"metadata": {},
"outputs": [],
"source": [
+ "%load_ext juliacall\n",
"%julia using Primes: prime"
]
},
@@ -1494,7 +1495,7 @@
"id": "S5dO61g1bDhk"
},
"source": [
- "The full list of PySR parameters can be found here: https://ai.damtp.cam.ac.uk/pysr/api"
+ "The full list of PySR parameters can be found here: https://pysr.ai/api"
]
}
],
diff --git a/examples/search_trace.py b/examples/search_trace.py
new file mode 100644
index 000000000..84228668f
--- /dev/null
+++ b/examples/search_trace.py
@@ -0,0 +1,159 @@
+"""Record the genealogy of a search: every expression evaluated, and its parents.
+
+`use_tracing` makes SymbolicRegression.jl write one JSONL record per (iteration, island)
+holding the live population and every mutation, crossover, tuning and death event. Joining
+those events on member refs gives the full parent-to-child graph, from which the winner's
+ancestry is a walk back up the parent edges.
+
+The trace writer lives in a JSON3 weak-dependency extension, which PySR installs into its
+Julia environment the first time a search asks for it.
+
+Trace records name features `x1`, `x2`, ... in column order, so `x1` below is `X[:, 0]`.
+"""
+
+import collections
+import json
+import pathlib
+import sys
+
+import numpy as np
+
+from pysr import PySRRegressor
+
+rng = np.random.default_rng(0)
+X = rng.uniform(-3, 3, size=(200, 3))
+y = 2.5382 * np.cos(X[:, 2]) + X[:, 0] ** 2 - 1.5
+y = y + rng.normal(0, 0.05 * y.std(), size=200)
+
+TRACE_PATH = "search_trace.jsonl"
+NITERATIONS = 25
+
+SETTINGS = dict(
+ binary_operators=["+", "-", "*", "/"],
+ unary_operators=["cos", "exp"],
+ population_size=27,
+ populations=8,
+ ncycles_per_iteration=90,
+ maxsize=20,
+ random_state=0,
+)
+
+
+def run_search(trace_path, niterations=NITERATIONS):
+ """Run the search with tracing on, leaving the JSONL at `trace_path`."""
+ model = PySRRegressor(
+ niterations=niterations,
+ use_tracing=True,
+ tracing_file=trace_path,
+ **SETTINGS,
+ )
+ model.fit(X, y)
+ return model
+
+
+def load_trace(trace_path):
+ """Read the JSONL into `nodes` (ref -> expression record) and `edges` (parent to child).
+
+ A crossover event is pushed onto both parents' event lists, so it is read from
+ `parent1` only. Each child records which of the two it descends from, so the edge is
+ oriented onto that parent and the other is carried alongside it as `mate`.
+ """
+ nodes, edges = {}, []
+ for line in pathlib.Path(trace_path).read_text().splitlines():
+ record = json.loads(line)
+ if record["record_type"] != "iteration":
+ continue
+ for member in record["members"]:
+ nodes.setdefault(member["ref"], {}).update(member)
+ for ref, entry in record["mutations"].items():
+ parent = int(ref)
+ nodes.setdefault(parent, {}).update(
+ {k: entry[k] for k in ("tree", "loss", "cost", "parent") if k in entry}
+ )
+ for event in entry["events"]:
+ kind = event["type"]
+ detail = event.get("mutation") or event.get("details") or {}
+ if kind == "crossover" and event["parent1"] != parent:
+ continue
+ children = (
+ (event["child1"], event["child2"])
+ if kind == "crossover"
+ else (event["child"],) if "child" in event else ()
+ )
+ for child in children:
+ edges.append(
+ dict(
+ parent=parent,
+ child=child,
+ kind=kind,
+ mutation=detail.get("type"),
+ result=detail.get("result", "accept"),
+ mate=event.get("parent2"),
+ )
+ )
+ for edge in edges:
+ if (
+ edge["mate"] is not None
+ and nodes.get(edge["child"], {}).get("parent") == edge["mate"]
+ ):
+ edge["parent"], edge["mate"] = edge["mate"], edge["parent"]
+ return nodes, edges
+
+
+def ancestry(nodes, edges):
+ """The lit path: the lowest-loss expression walked back up its parent edges."""
+ first_edge = {}
+ for edge in edges:
+ first_edge.setdefault(edge["child"], edge)
+ winner = min(
+ (ref for ref, n in nodes.items() if n.get("loss") is not None),
+ key=lambda ref: nodes[ref]["loss"],
+ )
+ chain, current, seen = [], winner, set()
+ while current in first_edge and current not in seen:
+ seen.add(current)
+ chain.append(first_edge[current])
+ current = first_edge[current]["parent"]
+ chain.reverse()
+ return winner, chain
+
+
+def check(nodes, edges):
+ """A usable genealogy: every edge resolves, the winner chains back to a seed, and
+ crossovers carry a distinct second parent."""
+ if not edges or not any(e["kind"] == "crossover" for e in edges):
+ return False
+ if not any(e["mate"] is not None and e["mate"] != e["parent"] for e in edges):
+ return False
+ if any(e["parent"] not in nodes for e in edges):
+ return False
+ winner, chain = ancestry(nodes, edges)
+ return bool(chain) and nodes[chain[0]["parent"]].get("parent") == -1
+
+
+def main():
+ trace_path = sys.argv[1] if len(sys.argv) > 1 else TRACE_PATH
+ run_search(trace_path)
+ nodes, edges = load_trace(trace_path)
+
+ kinds = collections.Counter(e["kind"] for e in edges)
+ print(f"trace file: {trace_path} ({pathlib.Path(trace_path).stat().st_size} bytes)")
+ print(f"expressions recorded: {len(nodes)}")
+ print(f"parent-to-child edges: {len(edges)} {dict(kinds)}")
+ print(f"rejected: {sum(1 for e in edges if e['result'] == 'reject')}")
+ print(
+ f"edges with a second parent: {sum(1 for e in edges if e['mate'] is not None and e['mate'] != e['parent'])}"
+ )
+
+ winner, chain = ancestry(nodes, edges)
+ print(f"\nwinner loss {nodes[winner]['loss']:.6g}, ancestry {len(chain)} steps")
+ for edge in chain[-8:]:
+ print(
+ f" {edge['kind']:10s} {str(edge['mutation']):32s} -> {nodes[edge['child']].get('tree', '?')[:70]}"
+ )
+
+ print(f"\nusable genealogy: {check(nodes, edges)}")
+
+
+if __name__ == "__main__":
+ main()
diff --git a/examples/test_examples.py b/examples/test_examples.py
new file mode 100644
index 000000000..a4c431c3e
--- /dev/null
+++ b/examples/test_examples.py
@@ -0,0 +1,274 @@
+"""Executable coverage for the shipped documentation examples.
+
+Every example in `examples/` is a plain script a reader can run directly. Most also
+expose the handful of names this suite needs, so the documented result is checked by
+running the documented search rather than by re-describing it here:
+
+ X, y the problem
+ VARIABLE_NAMES list[str], or absent
+ MODEL_KWARGS every PySRRegressor argument except random_state
+ check(model) -> bool the example's own definition of success
+
+Examples with their own runners are registered separately below.
+
+Wall times are single-seed measurements on a `genx` node, one thread, PySR 2.1.0.
+They pick the marker, not an assertion: timing a search on unknown CI hardware would
+be a flaky test.
+"""
+
+import ast
+import importlib.util
+import re
+import sys
+from pathlib import Path
+
+import pytest
+
+EXAMPLES_DIR = Path(__file__).resolve().parent
+
+# name -> slowest measured single-seed wall time in seconds
+EXAMPLES = {
+ "adaptive_mutation_weights": 28,
+ "affine_cipher": 5,
+ "game_of_life": 73,
+ "automatic_batching": 90,
+ "backsolve_mutation": 179,
+ "morton_interleave": 222,
+ "mutations_and_plugins": 214,
+ "any_arity": 362,
+ "turtle_graphics": 2414,
+ "cartpole_objective": 6866,
+ "prng_period": 3230,
+ "magnetic_field": 5507,
+}
+
+# Examples that cannot go through a single PySRRegressor fit and carry their own runner.
+SELF_DRIVEN = {"search_trace", "agent_loop_guesses", "hopper_controller"}
+
+# Examples whose search recovers its target on some seeds and not others. The measured
+# rate is reported in the docs; asserting recovery on one fixed seed would bake in a
+# lucky draw, so these are checked for completing and producing a usable front.
+STOCHASTIC = {"affine_freemod": 190}
+
+FAST_SECONDS = 100
+
+
+def load(name):
+ path = EXAMPLES_DIR / f"{name}.py"
+ spec = importlib.util.spec_from_file_location(name, path)
+ module = importlib.util.module_from_spec(spec)
+ sys.modules[name] = module
+ spec.loader.exec_module(module)
+ return module
+
+
+def run(name):
+ from pysr import PySRRegressor
+
+ module = load(name)
+ model = PySRRegressor(**module.MODEL_KWARGS, random_state=0)
+
+ fit_kwargs = {}
+ if getattr(module, "VARIABLE_NAMES", None):
+ fit_kwargs["variable_names"] = module.VARIABLE_NAMES
+ model.fit(module.X, module.y, **fit_kwargs)
+
+ assert module.check(model), f"{name} no longer reproduces its documented result"
+
+
+@pytest.mark.parametrize("name", [n for n, s in EXAMPLES.items() if s <= FAST_SECONDS])
+def test_example(name):
+ run(name)
+
+
+@pytest.mark.slow
+@pytest.mark.parametrize("name", [n for n, s in EXAMPLES.items() if s > FAST_SECONDS])
+def test_example_slow(name):
+ run(name)
+
+
+@pytest.mark.slow
+def test_agent_loop_reaches_floor():
+ """The loop's premise: a cold search misses the floor and fed-back guesses reach it."""
+ module = load("agent_loop_guesses")
+ rounds = module.main(seed=0)
+
+ assert len(rounds) > 1, "a cold start alone succeeded; the example has no loop left"
+ assert rounds[0]["loss"] > 2 * module.NOISE_FLOOR, "round 0 already at the floor"
+ assert (
+ rounds[-1]["loss"] <= 2 * module.NOISE_FLOOR
+ ), "the loop never reached the floor"
+
+
+@pytest.mark.slow
+def test_search_trace_genealogy(tmp_path):
+ """The trace resolves into a graph the winner's ancestry can be walked back through.
+
+ `use_tracing` pulls JSON3 into the PySR Julia environment on first use, so a cold
+ environment pays an install here before the search starts.
+ """
+ module = load("search_trace")
+ trace_path = str(tmp_path / "search_trace.jsonl")
+
+ module.run_search(trace_path, niterations=3)
+ nodes, edges = module.load_trace(trace_path)
+
+ assert module.check(nodes, edges), "the recorded genealogy no longer resolves"
+
+
+@pytest.mark.slow
+@pytest.mark.parametrize("name", sorted(STOCHASTIC))
+def test_stochastic_example_completes(name):
+ """The search runs and yields a front, whether or not this seed finds the answer."""
+ from pysr import PySRRegressor
+
+ module = load(name)
+ model = PySRRegressor(**module.MODEL_KWARGS, random_state=0)
+ model.fit(module.X, module.y)
+
+ assert not model.equations_.empty, f"{name} produced no front"
+ assert len(module.decode(model)) == len(module.PLAIN)
+
+
+def test_every_example_is_registered():
+ """A new file in examples/ must be registered, or it ships with no coverage."""
+ harness = {Path(__file__).stem, "conftest"}
+ on_disk = {p.stem for p in EXAMPLES_DIR.glob("*.py")} - harness
+ known = set(EXAMPLES) | SELF_DRIVEN | set(STOCHASTIC)
+ assert (
+ on_disk == known
+ ), f"unregistered: {sorted(on_disk - known)}, missing from disk: {sorted(known - on_disk)}"
+
+
+# Examples are ordered by what a reader needs first, not by which part of PySR they
+# exercise: the core of a search, then the knobs that shape one, then value types and
+# the searches built on them. Reordering is a deliberate act, so this list, the index
+# page, and the VitePress sidebar move together.
+READING_ORDER = [
+ "getting-started",
+ "expression-specifications",
+ "objectives",
+ "physics",
+ "search-behaviour",
+ "instrumentation",
+ "value-types",
+ "beyond-numeric-values",
+]
+
+DOCS_SRC = EXAMPLES_DIR.parent / "docs" / "src"
+
+
+def _squash(text):
+ return re.sub(r"[^a-z0-9]", "", text.lower())
+
+
+def test_docs_examples_follow_one_declared_order():
+ """Sidebar, index, and page bodies agree on the reading order, core material first."""
+ if not DOCS_SRC.exists():
+ pytest.skip("docs tree absent: examples were copied out of the repo")
+
+ config = (DOCS_SRC / ".vitepress" / "config.mts").read_text()
+ assert re.findall(r"link: '/examples/([a-z-]+)'", config) == READING_ORDER
+
+ index = (DOCS_SRC / "examples.md").read_text()
+ assert (
+ re.findall(r"^## \[[^]]+\]\(/examples/([a-z-]+)\)$", index, re.M)
+ == READING_ORDER
+ )
+
+ groups = dict(zip(READING_ORDER, index.split("\n## ")[1:]))
+ for slug in READING_ORDER:
+ listed = re.findall(
+ rf"^- \[[^]]+\]\(/examples/{slug}#([a-z0-9-]+)\)", groups[slug], re.M
+ )
+ page = (DOCS_SRC / "examples" / f"{slug}.md").read_text()
+ headings = [h for h in re.findall(r"^## (.+)$", page, re.M) if h != "Preamble"]
+ assert [_squash(a) for a in listed] == [
+ _squash(h) for h in headings
+ ], f"{slug}: index lists {listed}, page has {headings}"
+
+
+def _split_literals(source):
+ """Literals written as adjacent parts across lines, which hide their own seams."""
+ found = []
+ for node in ast.walk(ast.parse(source)):
+ if isinstance(node, ast.JoinedStr) or (
+ isinstance(node, ast.Constant) and isinstance(node.value, str)
+ ):
+ if node.lineno == node.end_lineno:
+ continue
+ segment = ast.get_source_segment(source, node) or ""
+ if not re.match(r"^[fFrRbB]{0,2}(\"\"\"|''')", segment):
+ found.append(segment.splitlines()[0].strip())
+ return found
+
+
+def test_no_string_is_split_across_lines():
+ """A string spanning lines is triple-quoted, never adjacent parts.
+
+ Implicit concatenation hides whitespace at the seams. These files are Julia source
+ a reader copies, where a dropped space silently changes the program.
+ """
+ offenders = {}
+ for path in sorted(EXAMPLES_DIR.glob("*.py")):
+ split = _split_literals(path.read_text())
+ if split:
+ offenders[path.name] = split
+ for page in sorted((DOCS_SRC / "examples").glob("*.md")):
+ for block in re.findall(
+ r"^```python\n(.*?)^```", page.read_text(), re.S | re.M
+ ):
+ try:
+ split = _split_literals(block)
+ except SyntaxError:
+ continue
+ if split:
+ offenders.setdefault(page.name, []).extend(split)
+ assert not offenders, f"strings split across lines: {offenders}"
+
+
+def _stray_indents(source):
+ """Multi-line literals whose body resets to a shallower column than its own opener."""
+ lines = source.split("\n")
+ found = []
+ for node in ast.walk(ast.parse(source)):
+ if not isinstance(node, (ast.JoinedStr, ast.Constant)):
+ continue
+ if isinstance(node, ast.Constant) and not isinstance(node.value, str):
+ continue
+ if node.lineno == node.end_lineno:
+ continue
+ opener = lines[node.lineno - 1]
+ body = [line for line in lines[node.lineno : node.end_lineno] if line.strip()]
+ if not body:
+ continue
+ base = len(opener) - len(opener.lstrip(" "))
+ shallowest = min(len(line) - len(line.lstrip(" ")) for line in body)
+ if shallowest != base:
+ found.append(opener.strip())
+ return found
+
+
+def test_embedded_julia_is_laid_out_like_its_host():
+ """A Julia block inside a Python file is indented with the Python around it.
+
+ Julia ignores indentation, so an `end` dragged out to column 0 buys nothing and
+ leaves the Python file looking broken at every closer.
+ """
+ offenders = {}
+ for path in sorted(EXAMPLES_DIR.glob("*.py")):
+ stray = _stray_indents(path.read_text())
+ if stray:
+ offenders[path.name] = stray
+ if DOCS_SRC.exists():
+ for page in sorted((DOCS_SRC / "examples").glob("*.md")):
+ for block in re.findall(
+ r"^```python\n(.*?)^```", page.read_text(), re.S | re.M
+ ):
+ try:
+ stray = _stray_indents(block)
+ except SyntaxError:
+ continue
+ if stray:
+ offenders.setdefault(page.name, []).extend(stray)
+ assert not offenders, f"literals dedented below their opener: {offenders}"
diff --git a/examples/turtle_graphics.py b/examples/turtle_graphics.py
new file mode 100644
index 000000000..7335100bf
--- /dev/null
+++ b/examples/turtle_graphics.py
@@ -0,0 +1,355 @@
+"""Search for a turtle program that draws a plus-sign outline.
+
+The value the search carries is a drawing: a list of pen commands, `forward` and
+`turn`. The operators move and turn the pen, so every expression is a program and
+its value is the picture that program draws. The loss rasterizes both drawings
+and compares them with a chamfer distance.
+"""
+
+import math
+
+import numpy as np
+
+from pysr import PySRRegressor, TypeSpec
+
+# --------------------------------------------------------------- the turtle
+# A Path is a flat command list carried as two parallel vectors: `kind` (0
+# forward, 1 turn) and `val` (a length, or an angle in radians).
+PREAMBLE = r"""
+import Random
+
+const _GRID = 96
+const _MARGIN = 3
+# Off-centre by an arbitrary fraction of a pixel: axis-aligned figures otherwise
+# land exactly on rounding ties, where a one-ulp change of an angle flips a pixel
+# and the loss stops being reproducible across implementations.
+const _OFFSET = 0.1373
+const _MAXCMD = 8192
+const _FAIL = 4.0
+const _DTCACHE = Dict{UInt64,Tuple{BitMatrix,Matrix{Float64}}}()
+const _DTLOCK = ReentrantLock()
+
+_empty() = (Int8[], Float64[])
+
+# ----------------------------------------------------- constants and printing
+# The width of the mutation step at temperature 1, in degrees.
+const _SIGMA_DEG = 30.0
+
+# One turn, any real angle in (-180, 180].
+_real(rng) = deg2rad(180.0 - 360.0 * rand(rng))
+
+function _wrap(deg)
+ d = mod(deg + 180.0, 360.0) - 180.0
+ return d <= -180.0 ? 180.0 : d
+end
+
+_perturb(rng, angle, temperature) =
+ deg2rad(_wrap(rad2deg(angle) + _SIGMA_DEG * temperature * Random.randn(rng)))
+
+# The optimizer sees exactly the turns: one Float64 per turn, forwards untouched.
+_turns(kind::Vector{Int8}, val::Vector{Float64}) =
+ [val[i] for i in eachindex(kind) if kind[i] == Int8(1)]
+
+function _put_turns(kind::Vector{Int8}, val::Vector{Float64}, c)
+ out = copy(val)
+ j = 0
+ for i in eachindex(kind)
+ if kind[i] == Int8(1)
+ j += 1
+ out[i] = c[j]
+ end
+ end
+ return out
+end
+
+function _degstr(angle)
+ d = rad2deg(angle)
+ r = round(d)
+ return abs(d - r) < 1e-9 ? string(Int(r)) : string(d)
+end
+
+_cmd(k::Int8, v::Float64) =
+ k == Int8(1) ? "t" * _degstr(v) : (v == 1.0 ? "f" : "f" * string(v))
+
+# Constant folding turns an all-constant subtree into one multi-command Path, so
+# a constant has to print as something that can be read back command by command.
+function _literal(kind::Vector{Int8}, val::Vector{Float64})
+ length(kind) == 1 && return _cmd(kind[1], val[1])
+ return "p[" * join([_cmd(kind[i], val[i]) for i in eachindex(kind)], ",") * "]"
+end
+
+# --------------------------------------------------------- the path algebra
+function _endpoint(kind::Vector{Int8}, val::Vector{Float64})
+ x = 0.0; y = 0.0; c = 1.0; s = 0.0
+ @inbounds for i in eachindex(kind)
+ if kind[i] == Int8(0)
+ x += c * val[i]; y += s * val[i]
+ else
+ ca = cos(val[i]); sa = sin(val[i])
+ c, s = c * ca - s * sa, s * ca + c * sa
+ end
+ end
+ return x, y
+end
+
+function _seq(ak::Vector{Int8}, av::Vector{Float64},
+ bk::Vector{Int8}, bv::Vector{Float64})
+ (isempty(ak) && isempty(bk)) && return _empty()
+ length(ak) + length(bk) > _MAXCMD && return _empty()
+ return (vcat(ak, bk), vcat(av, bv))
+end
+
+function _mir(kind::Vector{Int8}, val::Vector{Float64})
+ isempty(kind) && return _empty()
+ out = copy(val)
+ @inbounds for i in eachindex(kind)
+ kind[i] == Int8(1) && (out[i] = -out[i])
+ end
+ return (copy(kind), out)
+end
+
+# L-system self-application: replace every forward by a copy of the whole path,
+# scaled so the copy's net displacement equals that forward.
+function _nest(kind::Vector{Int8}, val::Vector{Float64})
+ isempty(kind) && return _empty()
+ nf = count(==(Int8(0)), kind)
+ nf == 0 && return _empty()
+ n = length(kind)
+ n + nf * (n - 1) > _MAXCMD && return _empty()
+ x, y = _endpoint(kind, val)
+ d = hypot(x, y)
+ d < 1e-9 && return _empty()
+ ok = Vector{Int8}(); ov = Vector{Float64}()
+ sizehint!(ok, n + nf * (n - 1)); sizehint!(ov, n + nf * (n - 1))
+ @inbounds for i in eachindex(kind)
+ if kind[i] == Int8(0)
+ r = val[i] / d
+ for j in eachindex(kind)
+ push!(ok, kind[j])
+ push!(ov, kind[j] == Int8(0) ? val[j] * r : val[j])
+ end
+ else
+ push!(ok, kind[i]); push!(ov, val[i])
+ end
+ end
+ return (ok, ov)
+end
+
+# ------------------------------------------------------------------ the loss
+# Walk the commands and stamp every drawn segment onto a 96x96 grid, after
+# centring the bounding box and scaling its larger extent to the grid, so the
+# comparison sees shape and not place or size.
+function _raster(kind::Vector{Int8}, val::Vector{Float64})
+ isempty(kind) && return nothing
+ nf = count(==(Int8(0)), kind)
+ nf == 0 && return nothing
+ xs = Vector{Float64}(undef, 4 * nf)
+ x = 0.0; y = 0.0; c = 1.0; s = 0.0
+ j = 0
+ xmin = Inf; xmax = -Inf; ymin = Inf; ymax = -Inf
+ @inbounds for i in eachindex(kind)
+ if kind[i] == Int8(0)
+ x2 = x + c * val[i]; y2 = y + s * val[i]
+ xs[4j + 1] = x; xs[4j + 2] = y; xs[4j + 3] = x2; xs[4j + 4] = y2
+ j += 1
+ xmin = min(xmin, x, x2); xmax = max(xmax, x, x2)
+ ymin = min(ymin, y, y2); ymax = max(ymax, y, y2)
+ x = x2; y = y2
+ else
+ ca = cos(val[i]); sa = sin(val[i])
+ c, s = c * ca - s * sa, s * ca + c * sa
+ end
+ end
+ (isfinite(xmin) && isfinite(xmax) && isfinite(ymin) && isfinite(ymax)) ||
+ return nothing
+ scale = max(xmax - xmin, ymax - ymin)
+ (scale < 1e-9 || !isfinite(scale)) && return nothing
+ span = _GRID - 1 - 2 * _MARGIN
+ cx = (xmin + xmax) / 2; cy = (ymin + ymax) / 2
+ mask = falses(_GRID, _GRID)
+ @inbounds for k in 0:(nf - 1)
+ u0 = _MARGIN + _OFFSET + ((xs[4k + 1] - cx) / scale + 0.5) * span
+ v0 = _MARGIN + _OFFSET + ((xs[4k + 2] - cy) / scale + 0.5) * span
+ u1 = _MARGIN + _OFFSET + ((xs[4k + 3] - cx) / scale + 0.5) * span
+ v1 = _MARGIN + _OFFSET + ((xs[4k + 4] - cy) / scale + 0.5) * span
+ steps = max(1, ceil(Int, 2 * hypot(u1 - u0, v1 - v0)))
+ for t in 0:steps
+ a = t / steps
+ i = clamp(round(Int, u0 + a * (u1 - u0)) + 1, 1, _GRID)
+ jj = clamp(round(Int, v0 + a * (v1 - v0)) + 1, 1, _GRID)
+ mask[i, jj] = true
+ end
+ end
+ return mask
+end
+
+# Exact squared Euclidean distance transform, one dimension at a time.
+function _dt1d!(d::Vector{Float64}, f::Vector{Float64}, n::Int,
+ v::Vector{Int}, z::Vector{Float64})
+ k = 1; v[1] = 1; z[1] = -1e20; z[2] = 1e20
+ @inbounds for q in 2:n
+ s = ((f[q] + q * q) - (f[v[k]] + v[k] * v[k])) / (2q - 2v[k])
+ while s <= z[k]
+ k -= 1
+ s = ((f[q] + q * q) - (f[v[k]] + v[k] * v[k])) / (2q - 2v[k])
+ end
+ k += 1; v[k] = q; z[k] = s; z[k + 1] = 1e20
+ end
+ k = 1
+ @inbounds for q in 1:n
+ while z[k + 1] < q
+ k += 1
+ end
+ d[q] = (q - v[k])^2 + f[v[k]]
+ end
+ return d
+end
+
+function _edt(mask::BitMatrix)
+ n1, n2 = size(mask)
+ big = 1e12
+ d2 = Matrix{Float64}(undef, n1, n2)
+ @inbounds for j in 1:n2, i in 1:n1
+ d2[i, j] = mask[i, j] ? 0.0 : big
+ end
+ m = max(n1, n2)
+ f = Vector{Float64}(undef, m); d = Vector{Float64}(undef, m)
+ v = Vector{Int}(undef, m + 1); z = Vector{Float64}(undef, m + 2)
+ @inbounds for j in 1:n2
+ for i in 1:n1; f[i] = d2[i, j]; end
+ _dt1d!(d, f, n1, v, z)
+ for i in 1:n1; d2[i, j] = d[i]; end
+ end
+ @inbounds for i in 1:n1
+ for j in 1:n2; f[j] = d2[i, j]; end
+ _dt1d!(d, f, n2, v, z)
+ for j in 1:n2; d2[i, j] = d[j]; end
+ end
+ return d2
+end
+
+# The target never changes, so its raster and its transform are computed once.
+function _target_grids(kind::Vector{Int8}, val::Vector{Float64})
+ key = hash(kind, hash(val))
+ lock(_DTLOCK) do
+ get!(_DTCACHE, key) do
+ mask = _raster(kind, val)
+ mask === nothing ? (falses(_GRID, _GRID), fill(1e12, _GRID, _GRID)) :
+ (mask, _edt(mask))
+ end
+ end
+end
+
+function _mean_to(mask::BitMatrix, d2::Matrix{Float64})
+ total = 0.0; n = 0
+ @inbounds for i in eachindex(mask)
+ if mask[i]
+ total += sqrt(d2[i]); n += 1
+ end
+ end
+ return n == 0 ? nothing : total / n
+end
+
+# Symmetric chamfer distance in pixels, divided by the grid width. Zero means
+# the two drawings lit exactly the same pixels.
+function _chamfer(pk::Vector{Int8}, pv::Vector{Float64},
+ tk::Vector{Int8}, tv::Vector{Float64})
+ a = _raster(pk, pv)
+ a === nothing && return _FAIL
+ b, db = _target_grids(tk, tv)
+ any(b) || return _FAIL
+ da = _edt(a)
+ ab = _mean_to(a, db); ba = _mean_to(b, da)
+ (ab === nothing || ba === nothing) && return _FAIL
+ v = 0.5 * (ab + ba) / _GRID
+ return isfinite(v) ? v : _FAIL
+end
+"""
+
+PATH = TypeSpec(
+ "Path",
+ fields={"kind": "Vector{Int8}", "val": "Vector{Float64}"},
+ preamble=PREAMBLE,
+ # A constant is one turn, any real angle in (-180, 180].
+ sample="rng -> Path(Int8[1], [_real(rng)])",
+ # Perturb one of the value's angles by a temperature-scaled step, or
+ # resample it outright. Constant folding can hand this hook a value of
+ # several commands, so it moves a command in place.
+ mutate="""(rng, value, temperature) -> begin
+ turns = findall(isone, value.kind)
+ isempty(turns) && return Path(Int8[1], [_real(rng)])
+ i = rand(rng, turns)
+ val = copy(value.val)
+ val[i] = rand(rng, Bool) ? _perturb(rng, val[i], temperature) : _real(rng)
+ Path(copy(value.kind), val)
+ end""",
+ # Every turn in the value is a scalar the constant optimizer may fit.
+ scalar_constants="value -> _turns(value.kind, value.val)",
+ with_scalar_constants="(value, c) -> Path(copy(value.kind), _put_turns(value.kind, value.val, c))",
+ is_valid="value -> !isempty(value.kind) && all(isfinite, value.val)",
+ string="value -> _literal(value.kind, value.val)",
+)
+
+OPERATORS = {
+ 1: [
+ "dup(a::Path) = (t = _seq(a.kind, a.val, a.kind, a.val); Path(t[1], t[2]))",
+ "tri(a::Path) = (t = _seq(a.kind, a.val, a.kind, a.val); u = _seq(t[1], t[2], a.kind, a.val); Path(u[1], u[2]))",
+ "mir(a::Path) = (t = _mir(a.kind, a.val); Path(t[1], t[2]))",
+ "nest(a::Path) = (t = _nest(a.kind, a.val); Path(t[1], t[2]))",
+ ],
+ 2: [
+ "seq(a::Path, b::Path) = (t = _seq(a.kind, a.val, b.kind, b.val); Path(t[1], t[2]))"
+ ],
+}
+
+CHAMFER_LOSS = "chamfer(prediction::Path, target::Path)::Float64 = _chamfer(prediction.kind, prediction.val, target.kind, target.val)"
+
+
+def path(commands):
+ """A ('F', length) / ('T', degrees) command list as the two Path vectors."""
+ kind = [0 if op == "F" else 1 for op, _ in commands]
+ val = [a if op == "F" else math.radians(a) for op, a in commands]
+ return kind, val
+
+
+# The target: a plus-sign outline, four copies of a three-edge arm.
+ARM = [("F", 1.0), ("T", 90.0), ("F", 1.0), ("T", -90.0), ("F", 1.0), ("T", 90.0)]
+TARGET = ARM * 4
+
+# One row. The only feature is `f`, a single unit step forward.
+X = np.empty((1, 1), dtype=object)
+X[0, 0] = ([0], [1.0])
+y = np.empty(1, dtype=object)
+y[0] = path(TARGET)
+VARIABLE_NAMES = ["f"]
+
+MODEL_KWARGS = dict(
+ type_spec=PATH,
+ operators=OPERATORS,
+ elementwise_loss=CHAMFER_LOSS,
+ niterations=20,
+ verbosity=0,
+)
+
+EXACT = 1e-9
+
+
+def check(model):
+ """Some program on the front draws the target, to floating-point precision."""
+ return float(model.equations_["loss"].min()) < EXACT
+
+
+def main():
+ model = PySRRegressor(**MODEL_KWARGS, random_state=0)
+ model.fit(X, y, variable_names=VARIABLE_NAMES)
+ print(model.equations_[["complexity", "loss", "equation"]].to_string(index=False))
+ exact = model.equations_[model.equations_["loss"] < EXACT]
+ if len(exact):
+ best = exact.loc[exact["complexity"].idxmin()]
+ print(f"\ndraws the target: {best.equation}")
+ else:
+ print("\nno program drew the target")
+
+
+if __name__ == "__main__":
+ main()
diff --git a/pyproject.toml b/pyproject.toml
index fccd37f7a..a19729970 100644
--- a/pyproject.toml
+++ b/pyproject.toml
@@ -4,25 +4,25 @@ build-backend = "hatchling.build"
[project]
name = "pysr"
-version = "2.0.0a2"
+version = "2.6.0"
authors = [
{name = "Miles Cranmer", email = "miles.cranmer@gmail.com"},
]
description = "Simple and efficient symbolic regression"
readme = {file = "README.md", content-type = "text/markdown"}
license = {file = "LICENSE"}
-requires-python = ">=3.9"
+requires-python = ">=3.10"
classifiers = [
"Programming Language :: Python :: 3",
"Operating System :: OS Independent",
"License :: OSI Approved :: Apache Software License"
]
dependencies = [
- "sympy>=1.0.0,<2.0.0",
+ "sympy>=1.14.0,<2.0.0",
"pandas>=0.21.0,<4.0.0",
"numpy>=1.13.0,<3.0.0",
"scikit_learn>=1.0.0,<2.0.0",
- "juliacall>=0.9.28,<0.9.29",
+ "juliacall>=0.9.36,<0.9.37",
"click>=7.0.0,<9.0.0",
"typing-extensions>=4.0.0,<5.0.0",
]
@@ -35,18 +35,18 @@ docs = [
dev = [
"coverage>=7,<8",
"beartype>=0.19,<0.23",
- "ipykernel>=6,<7",
- "ipython>=8,<9",
- "jax[cpu]>=0.4,<0.6",
+ "ipykernel>=6,<8",
+ "ipython>=8,<10",
+ "jax[cpu]>=0.4,<0.12",
"jupyter>=1,<2",
"mypy>=1,<3",
"nbval>=0.11,<0.12",
"pandas-stubs",
"pre-commit>=3.0,<5",
- "pytest-cov>=5,<7",
+ "pytest-cov>=5,<8",
"pytest>=8,<10",
"tensorboard>=2,<3",
- "torch>=2,<3",
+ "torch>=2.7,<3",
"types-openpyxl",
"types-pytz",
]
diff --git a/pysr/__init__.py b/pysr/__init__.py
index 3e3f7c8be..8a1f5cb36 100644
--- a/pysr/__init__.py
+++ b/pysr/__init__.py
@@ -29,12 +29,35 @@
from .expression_specs import (
AbstractExpressionSpec,
ExpressionSpec,
- ParametricExpressionSpec,
TemplateExpressionSpec,
)
from .julia_extensions import load_all_packages
from .logger_specs import AbstractLoggerSpec, TensorBoardLoggerSpec
+from .mutations import (
+ AbstractMutation,
+ AddNodeMutation,
+ BacksolveMutation,
+ ConstantMutation,
+ DeleteNodeMutation,
+ DoNothingMutation,
+ FeatureMutation,
+ InsertNodeMutation,
+ OperatorMutation,
+ OptimizeMutation,
+ RandomizeMutation,
+ RotateTreeMutation,
+ SimplifyMutation,
+ SwapOperandsMutation,
+)
+from .plugins import (
+ AbstractPlugin,
+ AdaptiveMutationWeightsPlugin,
+ AdaptiveParsimonyPlugin,
+ MutationBurstPlugin,
+ SimulatedAnnealingPlugin,
+)
from .sr import PySRRegressor
+from .type_specs import TypeSpec
try:
__version__ = version("pysr")
@@ -54,9 +77,28 @@
"AbstractExpressionSpec",
"ExpressionSpec",
"TemplateExpressionSpec",
- "ParametricExpressionSpec",
"AbstractLoggerSpec",
"TensorBoardLoggerSpec",
+ "TypeSpec",
+ "AbstractMutation",
+ "ConstantMutation",
+ "OperatorMutation",
+ "FeatureMutation",
+ "SwapOperandsMutation",
+ "AddNodeMutation",
+ "InsertNodeMutation",
+ "DeleteNodeMutation",
+ "RotateTreeMutation",
+ "BacksolveMutation",
+ "SimplifyMutation",
+ "RandomizeMutation",
+ "OptimizeMutation",
+ "DoNothingMutation",
+ "AbstractPlugin",
+ "AdaptiveMutationWeightsPlugin",
+ "AdaptiveParsimonyPlugin",
+ "MutationBurstPlugin",
+ "SimulatedAnnealingPlugin",
"best",
"best_callable",
"best_row",
diff --git a/pysr/_cli/main.py b/pysr/_cli/main.py
index cd81710ff..f1a06fbca 100644
--- a/pysr/_cli/main.py
+++ b/pysr/_cli/main.py
@@ -1,4 +1,5 @@
import fnmatch
+import os
import sys
import unittest
import warnings
@@ -10,6 +11,7 @@
runtests,
runtests_autodiff,
runtests_dev,
+ runtests_interrupt,
runtests_jax,
runtests_slurm,
runtests_startup,
@@ -50,7 +52,17 @@ def _install(julia_project, quiet, precompile):
)
-TEST_OPTIONS = {"main", "jax", "torch", "autodiff", "cli", "dev", "startup", "slurm"}
+TEST_OPTIONS = {
+ "main",
+ "jax",
+ "torch",
+ "autodiff",
+ "cli",
+ "dev",
+ "startup",
+ "slurm",
+ "interrupt",
+}
@pysr.command("test")
@@ -65,9 +77,23 @@ def _install(julia_project, quiet, precompile):
def _tests(tests, expressions):
"""Run parts of the PySR test suite.
- Choose from main, jax, torch, autodiff, cli, dev, startup, and slurm.
+ Choose from main, jax, torch, autodiff, cli, dev, startup, slurm, and interrupt.
You can give multiple tests, separated by commas.
"""
+ try:
+ shard_count = int(os.environ.get("PYSR_TEST_SHARD_COUNT", "1"))
+ shard_index = int(os.environ.get("PYSR_TEST_SHARD_INDEX", "0"))
+ except ValueError as error:
+ raise click.ClickException(
+ "PYSR_TEST_SHARD_COUNT and PYSR_TEST_SHARD_INDEX must be integers."
+ ) from error
+ if shard_count <= 0:
+ raise click.ClickException("PYSR_TEST_SHARD_COUNT must be a positive integer.")
+ if not 0 <= shard_index < shard_count:
+ raise click.ClickException(
+ "PYSR_TEST_SHARD_INDEX must satisfy 0 <= PYSR_TEST_SHARD_INDEX "
+ "< PYSR_TEST_SHARD_COUNT."
+ )
test_cases = []
for test in tests.split(","):
if test == "main":
@@ -85,6 +111,8 @@ def _tests(tests, expressions):
test_cases.extend(runtests_dev(just_tests=True))
elif test == "startup":
test_cases.extend(runtests_startup(just_tests=True))
+ elif test == "interrupt":
+ test_cases.extend(runtests_interrupt(just_tests=True))
elif test == "slurm":
test_cases.extend(runtests_slurm(just_tests=True))
else:
@@ -100,6 +128,7 @@ def _tests(tests, expressions):
for expression in expressions
):
suite.addTest(test)
+ suite = unittest.TestSuite(list(suite)[shard_index::shard_count])
runner = unittest.TextTestRunner()
results = runner.run(suite)
diff --git a/pysr/expression_specs.py b/pysr/expression_specs.py
index bd8d5ec71..187516d5d 100644
--- a/pysr/expression_specs.py
+++ b/pysr/expression_specs.py
@@ -1,11 +1,10 @@
from __future__ import annotations
import copy
-import textwrap
-import warnings
from abc import ABC, abstractmethod
+from dataclasses import dataclass
from textwrap import dedent
-from typing import TYPE_CHECKING, Any, NewType, overload
+from typing import TYPE_CHECKING, Any, ClassVar, NewType
import numpy as np
import pandas as pd
@@ -13,7 +12,6 @@
from .export import add_export_formats
from .julia_helpers import jl_array
from .julia_import import AnyValue, SymbolicRegression, jl
-from .utils import ArrayLike
try:
from typing import TypeAlias
@@ -42,9 +40,7 @@ class AbstractExpressionSpec(ABC):
2. create_exports(), which will be used to create the exports of the equations, such as
the executable format, the SymPy format, etc.
- It may also optionally implement:
-
- - supports_sympy, supports_torch, supports_jax, supports_latex: Whether this expression type supports the corresponding export format.
+ Implementations may also declare the export formats they support.
"""
@abstractmethod
@@ -52,6 +48,30 @@ def julia_expression_spec(self) -> AnyValue:
"""The expression specification"""
pass # pragma: no cover
+ def _julia_expression_spec_source(self, *, prototype: str | None) -> str | None:
+ """Return self-contained Julia source for a TypeSpec-compatible spec.
+
+ ``prototype`` is Julia source for one value of the generated type. PySR
+ evaluates the returned source once in the fingerprinted runtime module.
+ """
+ return None
+
+ def _julia_expression_spec_function_selector(self) -> str | None:
+ """Return Julia source selecting the callable nested in the spec.
+
+ PySR evaluates the source to a selector, invokes it with the expression
+ spec, and binds the result under a deterministic runtime-module name so
+ checkpoints and workers resolve the same callable identity.
+ """
+ return None
+
+ @property
+ def supports_type_spec(self) -> bool:
+ return False
+
+ def _validate_type_spec(self) -> None:
+ pass
+
@abstractmethod
def create_exports(
self,
@@ -108,6 +128,10 @@ def create_exports(
output_torch_format=model.output_torch_format,
)
+ @property
+ def supports_type_spec(self) -> bool:
+ return True
+
@property
def supports_sympy(self):
return True
@@ -125,6 +149,7 @@ def supports_latex(self):
return True
+@dataclass
class TemplateExpressionSpec(AbstractExpressionSpec):
"""Spec for templated expressions.
@@ -177,211 +202,54 @@ class TemplateExpressionSpec(AbstractExpressionSpec):
the derivative of f with respect to its first argument, evaluated at x.
"""
- _spec_cache: dict[tuple[str, ...], AnyValue] = {}
+ combine: str
+ expressions: list[str]
+ variable_names: list[str]
+ parameters: dict[str, int] | None = None
- @overload
- def __init__(
- self,
- function_symbols: list[str],
- combine: str,
- num_features: dict[str, int] | None = None,
- ) -> None: ...
-
- @overload
- def __init__(
- self,
- combine: str,
- *,
- expressions: list[str],
- variable_names: list[str],
- parameters: dict[str, int] | None = None,
- ) -> None: ...
-
- def __init__(
- self,
- *args,
- **kwargs,
- ):
- """Handle both formats with combine as explicit parameter"""
- self._old_format = len(args) >= 2 or "function_symbols" in kwargs
-
- if self._old_format:
- self._load_old_format(*args, **kwargs)
- else:
- self._load_new_format(*args, **kwargs)
-
- def _load_old_format(
- self,
- function_symbols: list[str],
- combine: str,
- num_features: dict[str, int] | None = None,
- ):
- self.function_symbols = function_symbols
- self.combine = combine
- self.num_features = num_features
- # TODO: warn about old format after some versions
-
- def _load_new_format(
- self,
- combine: str,
- *,
- expressions: list[str],
- variable_names: list[str],
- parameters: dict[str, int] | None = None,
- ):
- self.combine = combine
- self.expressions = expressions
- self.variable_names = variable_names
- self.parameters = parameters
+ _spec_cache: ClassVar[dict[tuple[str, ...], AnyValue]] = {}
def _get_cache_key(self):
- if self._old_format:
- return (
- "old",
- str(self.function_symbols),
- self.combine,
- str(self.num_features),
- )
- else:
- return (
- "new",
- self.combine,
- str(self.expressions),
- str(self.variable_names),
- str(self.parameters),
- )
+ return (
+ self.combine,
+ str(self.expressions),
+ str(self.variable_names),
+ str(self.parameters),
+ )
def julia_expression_spec(self):
key = self._get_cache_key()
- if key in self._spec_cache:
- return self._spec_cache[key]
-
- if self._old_format:
- result = SymbolicRegression.TemplateExpressionSpec(
- structure=self.julia_expression_options().structure
- )
- else:
- result = self._call_template_macro()
-
- self._spec_cache[key] = result
- return result
+ if key not in self._spec_cache:
+ self._spec_cache[key] = self._call_template_macro()
+ return self._spec_cache[key]
def _call_template_macro(self):
return jl.seval(self._template_macro_str())
- def _template_macro_str(self):
+ def _template_macro_str(self, *, prototype: str | None = None) -> str:
template_inputs = [f"expressions=({', '.join(self.expressions) + ','})"]
if self.parameters:
template_inputs.append(
f"parameters=({', '.join([f'{p}={self.parameters[p]}' for p in self.parameters]) + ','})"
)
+ if prototype is not None:
+ template_inputs.append(f"prototype={prototype}")
return dedent(f"""
@template_spec({', '.join(template_inputs) + ','}) do {', '.join(self.variable_names)}
{self.combine}
end
""")
- def julia_expression_options(self):
- f_combine = jl.seval(self.combine)
- creator = jl.seval("""
- function _pysr_create_template_structure(
- @nospecialize(function_symbols::AbstractVector),
- @nospecialize(combine::Function),
- @nospecialize(num_features::Union{Nothing,AbstractDict})
- )
- tuple_symbol = (map(Symbol, function_symbols)..., )
- num_features = if num_features === nothing
- nothing
- else
- NamedTuple(Symbol(k) => v for (k, v) in num_features)
- end
- structure = SymbolicRegression.TemplateStructure{tuple_symbol}(combine, num_features)
- return (; structure)
- end
- """)
- return creator(self.function_symbols, f_combine, self.num_features)
+ def _julia_expression_spec_source(self, *, prototype: str | None) -> str:
+ return self._template_macro_str(prototype=prototype)
+
+ def _julia_expression_spec_function_selector(self) -> str:
+ return "spec -> spec.structure.combine"
@property
- def evaluates_in_julia(self):
+ def supports_type_spec(self) -> bool:
return True
- def create_exports(
- self,
- model: PySRRegressor,
- equations: pd.DataFrame,
- search_output,
- i: int | None = None,
- ) -> pd.DataFrame:
- # We try to load the raw julia state from a saved binary stream
- # if not provided.
- search_output = search_output or model.julia_state_
- return _search_output_to_callable_expressions(equations, search_output, i)
-
-
-def parametric_expression_deprecation_warning(
- max_parameters: int, variable_names: ArrayLike[str]
-):
- function_name = "f"
- var_names = list(variable_names)
- message = dedent(f"""
- ParametricExpressionSpec is deprecated. you should switch to TemplateExpressionSpec with explicit parameters indexed by category.
-
- Since you have `max_parameters={max_parameters}` and `variable_names=[{", ".join(f'"{v}"' for v in var_names)}]`, you could migrate like this:
-
- n_categories = len(np.unique(category)) # count the number of parameters required
- expression_spec = TemplateExpressionSpec(
- expressions=["{function_name}"],
- variable_names=[{", ".join(f'"{v}"' for v in var_names + ["category"])}],
- parameters={{{", ".join(f'"p{i+1}": n_categories' for i in range(max_parameters))}}},
- combine="{function_name}({', '.join(var_names + [f'p{i+1}[category]' for i in range(max_parameters)])})",
- )
- X = np.column_stack([X, category]) # add the category column
-
- Finally, do not pass `category` when calling .fit().
- """).strip()
- wrapped = "\n".join(textwrap.fill(line, 88) for line in message.splitlines())
- warnings.warn(wrapped, FutureWarning, stacklevel=3)
-
-
-class ParametricExpressionSpec(AbstractExpressionSpec):
- """Spec for parametric expressions that vary by category.
-
- **This is deprecated in favor of the `TemplateExpressionSpec` class,
- which now supports parameters indexed by category.**
-
- This class allows you to specify expressions with parameters that vary across different
- categories in your dataset. The expression structure remains the same, but parameters
- are optimized separately for each category.
-
- Parameters
- ----------
- max_parameters : int
- Maximum number of parameters that can appear in the expression. Each parameter
- will take on different values for each category in the data.
-
- Examples
- --------
- For example, if we want to allow for a model with up to 2 parameters (each category
- can have a different value for these parameters), we can use:
-
- ```python
- model = PySRRegressor(
- expression_spec=ParametricExpressionSpec(max_parameters=2),
- binary_operators=["+", "*"],
- unary_operators=["sin"]
- )
- model.fit(X, y, category=category)
- ```
- """
-
- def __init__(self, max_parameters: int):
- self.max_parameters = max_parameters
-
- def julia_expression_spec(self):
- return SymbolicRegression.ParametricExpressionSpec(
- max_parameters=self.max_parameters, warn=False
- )
-
@property
def evaluates_in_julia(self):
return True
@@ -392,7 +260,7 @@ def create_exports(
equations: pd.DataFrame,
search_output,
i: int | None = None,
- ):
+ ) -> pd.DataFrame:
search_output = search_output or model.julia_state_
return _search_output_to_callable_expressions(equations, search_output, i)
@@ -407,7 +275,9 @@ def __call__(self, X: np.ndarray, *args):
def _search_output_to_callable_expressions(
- equations: pd.DataFrame, search_output, i: int | None
+ equations: pd.DataFrame,
+ search_output,
+ i: int | None,
) -> pd.DataFrame:
equations = copy.deepcopy(equations)
_, all_out_hof = search_output
diff --git a/pysr/interrupt.py b/pysr/interrupt.py
new file mode 100644
index 000000000..cd1d9b4ad
--- /dev/null
+++ b/pysr/interrupt.py
@@ -0,0 +1,132 @@
+"""Cooperative SIGINT handling so searches stop gracefully instead of killing the process."""
+
+from __future__ import annotations
+
+import _thread
+import ctypes
+import os
+import signal
+import socket
+import threading
+import warnings
+from contextlib import ExitStack, contextmanager
+from ctypes import wintypes
+
+from .julia_import import SymbolicRegression
+
+
+class _SigactionStorage(ctypes.Union):
+ _fields_ = [
+ ("alignment", ctypes.c_longdouble),
+ ("storage", ctypes.c_ubyte * 1024),
+ ]
+
+
+def _libc_with_sigaction():
+ libc = ctypes.CDLL(None, use_errno=True)
+ libc.sigaction.argtypes = [ctypes.c_int, ctypes.c_void_p, ctypes.c_void_p]
+ libc.sigaction.restype = ctypes.c_int
+ return libc
+
+
+def _checked_sigaction(libc, action, old_action):
+ result = libc.sigaction(signal.SIGINT, action, old_action)
+ if result != 0:
+ errno = ctypes.get_errno()
+ raise OSError(errno, os.strerror(errno))
+
+
+@contextmanager
+def _console_ctrl_c_to_python():
+ """Hand console Ctrl-C to Python before Julia's handler exits the process.
+
+ Windows calls console handlers last-registered first, so this one runs first.
+ """
+ kernel32 = ctypes.WinDLL("kernel32", use_last_error=True)
+ HANDLER = ctypes.WINFUNCTYPE(wintypes.BOOL, wintypes.DWORD)
+ kernel32.SetConsoleCtrlHandler.argtypes = [HANDLER, wintypes.BOOL]
+ kernel32.SetConsoleCtrlHandler.restype = wintypes.BOOL
+
+ def handler(event):
+ if event == 0: # CTRL_C_EVENT
+ _thread.interrupt_main()
+ return True
+ return False
+
+ callback = HANDLER(handler)
+ if not kernel32.SetConsoleCtrlHandler(callback, True):
+ raise ctypes.WinError(ctypes.get_last_error())
+ try:
+ yield
+ finally:
+ if not kernel32.SetConsoleCtrlHandler(callback, False):
+ raise ctypes.WinError(ctypes.get_last_error())
+
+
+def _should_arm_external_stop() -> bool:
+ return (
+ threading.current_thread() is threading.main_thread()
+ and os.environ.get("PYTHON_JULIACALL_HANDLE_SIGNALS") == "yes"
+ )
+
+
+def _stop_channel(cleanup: ExitStack) -> tuple[int, int]:
+ """Open a nonblocking pair of descriptors carrying the signal wakeup byte.
+
+ Windows writes that byte with `send` and can only poll sockets, so a socket
+ pair is the one transport that works everywhere.
+ """
+ read_end, write_end = socket.socketpair()
+ cleanup.callback(read_end.close)
+ cleanup.callback(write_end.close)
+ read_end.setblocking(False)
+ write_end.setblocking(False)
+ return read_end.fileno(), write_end.fileno()
+
+
+@contextmanager
+def _external_stop_signal_context(model):
+ interrupted = False
+ external_stop = None
+
+ with ExitStack() as cleanup:
+ if _should_arm_external_stop():
+ stop_read_fd, stop_write_fd = _stop_channel(cleanup)
+ external_stop = SymbolicRegression.ExternalStop(stop_read_fd, signal.SIGINT)
+
+ # `signal.signal` below replaces the SIGINT handler Julia installed
+ # on POSIX, and sigaction is the only way to put it back.
+ if os.name == "posix":
+ libc = _libc_with_sigaction()
+ saved_sigaction = _SigactionStorage()
+ _checked_sigaction(libc, None, ctypes.byref(saved_sigaction))
+ cleanup.callback(
+ _checked_sigaction, libc, ctypes.byref(saved_sigaction), None
+ )
+
+ saved_python_handler = signal.getsignal(signal.SIGINT)
+
+ def record_interrupt(*_):
+ nonlocal interrupted
+ interrupted = True
+
+ signal.signal(signal.SIGINT, record_interrupt)
+ cleanup.callback(signal.signal, signal.SIGINT, saved_python_handler)
+ previous_wakeup_fd = signal.set_wakeup_fd(stop_write_fd)
+ cleanup.callback(signal.set_wakeup_fd, previous_wakeup_fd)
+
+ # On Windows, Julia takes Ctrl-C through a console handler instead.
+ # Registering ours last means it is removed first, while
+ # `record_interrupt` still absorbs a Ctrl-C arriving mid-teardown.
+ if os.name != "posix":
+ cleanup.enter_context(_console_ctrl_c_to_python())
+
+ yield external_stop
+
+ model.interrupted_ = interrupted
+ if interrupted:
+ warnings.warn(
+ "The search was interrupted. Returning partial results.",
+ RuntimeWarning,
+ stacklevel=3,
+ )
diff --git a/pysr/julia_extensions.py b/pysr/julia_extensions.py
index 6d16756e5..e9cd18200 100644
--- a/pysr/julia_extensions.py
+++ b/pysr/julia_extensions.py
@@ -16,6 +16,7 @@ def load_required_packages(
autodiff_backend: Literal["Zygote", "Mooncake", "Enzyme"] | None = None,
cluster_manager: str | None = None,
logger_spec: AbstractLoggerSpec | None = None,
+ use_tracing: bool = False,
):
if turbo:
load_package("LoopVectorization", "bdcacae8-1622-11e9-2a5c-532679323890")
@@ -27,10 +28,14 @@ def load_required_packages(
load_package("Mooncake", "da2b9cff-9c12-43a0-ae48-6db2b0edb7d6")
elif autodiff_backend == "Enzyme":
load_package("Enzyme", "7da242da-08ed-463a-9acd-ee780be4f1d9")
- if cluster_manager is not None:
+ if cluster_manager == "slurm":
+ load_package("SlurmClusterManager", "c82cd089-7bf7-41d7-976b-6b5d413cbe0a")
+ elif cluster_manager is not None:
load_package("ClusterManagers", "34f1f09b-3a8b-5176-ab39-66d58a4d544e")
if isinstance(logger_spec, TensorBoardLoggerSpec):
load_package("TensorBoardLogger", "899adc3e-224a-11e9-021f-63837185c80f")
+ if use_tracing:
+ load_package("JSON", "682c06a0-de6a-54ab-a142-c8b1cf79cde6")
def load_all_packages():
@@ -41,7 +46,9 @@ def load_all_packages():
autodiff_backend="Zygote",
cluster_manager="slurm",
logger_spec=TensorBoardLoggerSpec(log_dir="logs"),
+ use_tracing=True,
)
+ load_package("ClusterManagers", "34f1f09b-3a8b-5176-ab39-66d58a4d544e")
# TODO: Refactor this file so we can install all packages at once using `juliapkg`,
@@ -49,7 +56,9 @@ def load_all_packages():
def isinstalled(uuid_s: str):
- return jl.haskey(Pkg.dependencies(), jl.Base.UUID(uuid_s))
+ dependencies = Pkg.dependencies()
+ uuid = jl.Base.UUID(uuid_s)
+ return jl.haskey(dependencies, uuid) and dependencies[uuid].is_direct_dep
def load_package(package_name: str, uuid_s: str) -> None:
diff --git a/pysr/julia_helpers.py b/pysr/julia_helpers.py
index b72793105..f7411d1df 100644
--- a/pysr/julia_helpers.py
+++ b/pysr/julia_helpers.py
@@ -30,6 +30,19 @@ def _escape_filename(filename):
def _load_cluster_manager(cluster_manager: str):
+ if cluster_manager == "slurm":
+ jl.seval("using Distributed: addprocs")
+ jl.seval("using SlurmClusterManager: SlurmManager")
+ return jl.seval("""
+ (numprocs; kws...) -> begin
+ manager = SlurmManager()
+ manager.ntasks == numprocs || error(
+ "Requested $numprocs processes, but Slurm allocation has $(manager.ntasks) tasks. " *
+ "Set Slurm `--ntasks`/`--ntasks-per-node` and `procs` to the same value."
+ )
+ addprocs(manager; kws...)
+ end
+ """)
jl.seval(f"using ClusterManagers: addprocs_{cluster_manager}")
return jl.seval(f"addprocs_{cluster_manager}")
@@ -43,6 +56,10 @@ def jl_array(x, dtype=None):
return jl_convert(jl.Array[dtype], x)
+def jl_numpy_array(x):
+ return jl.copy(PythonCall.PyArray(x))
+
+
def jl_dict(x):
return jl_convert(jl.Dict, x)
@@ -52,7 +69,7 @@ def jl_named_tuple(d):
def jl_is_function(f) -> bool:
- return cast(bool, jl.seval("op -> op isa Function")(f))
+ return cast(bool, jl.isa(f, jl.Function))
def jl_serialize(obj: Any) -> NDArray[np.uint8]:
diff --git a/pysr/julia_import.py b/pysr/julia_import.py
index 0de6f499a..72c2610e4 100644
--- a/pysr/julia_import.py
+++ b/pysr/julia_import.py
@@ -6,6 +6,18 @@
from .julia_registry_helpers import try_with_registry_fallback
+autoload_extensions = os.environ.get("PYSR_AUTOLOAD_EXTENSIONS")
+if autoload_extensions is not None:
+ warnings.warn(
+ "PYSR_AUTOLOAD_EXTENSIONS is deprecated and will be removed in a future release. "
+ "Set PYTHON_JULIACALL_AUTOLOAD_IPYTHON_EXTENSION instead.",
+ FutureWarning,
+ )
+ # Only fill the gap; an explicitly set juliacall variable takes precedence.
+ os.environ.setdefault(
+ "PYTHON_JULIACALL_AUTOLOAD_IPYTHON_EXTENSION", autoload_extensions
+ )
+
# Check if JuliaCall is already loaded, and if so, warn the user
# about the relevant environment variables. If not loaded,
# set up sensible defaults.
@@ -35,16 +47,18 @@
("PYTHON_JULIACALL_HANDLE_SIGNALS", "yes"),
("PYTHON_JULIACALL_THREADS", "auto"),
("PYTHON_JULIACALL_OPTLEVEL", "3"),
+ # Don't hijack `%%julia` magics in notebooks unless asked;
+ # opt back in with PYTHON_JULIACALL_AUTOLOAD_IPYTHON_EXTENSION=yes
+ # or `%load_ext juliacall`.
+ ("PYTHON_JULIACALL_AUTOLOAD_IPYTHON_EXTENSION", "no"),
+ # Concurrent GC sweeping (`N,1`): the search is GC-bound at high
+ # thread counts, and this was neutral at 8 threads, 8% faster at 32,
+ # and 25% faster at 96.
+ ("JULIA_NUM_GC_THREADS", f"{max(1, (os.cpu_count() or 2) // 2)},1"),
):
os.environ[k] = os.environ.get(k, default)
-autoload_extensions = os.environ.get("PYSR_AUTOLOAD_EXTENSIONS")
-if autoload_extensions is not None:
- # Deprecated; so just pass to juliacall
- os.environ["PYTHON_JULIACALL_AUTOLOAD_IPYTHON_EXTENSION"] = autoload_extensions
-
-
def _import_juliacall():
import juliacall # type: ignore
diff --git a/pysr/juliapkg.json b/pysr/juliapkg.json
index 532a94a21..f786583e2 100644
--- a/pysr/juliapkg.json
+++ b/pysr/juliapkg.json
@@ -3,8 +3,10 @@
"packages": {
"SymbolicRegression": {
"uuid": "8254be44-1295-4e6a-a16d-46603ac705cb",
- "url": "https://github.com/astroautomata/SymbolicRegression.jl",
- "rev": "v2.0.0-alpha.9"
+ "version": "~2.5.0",
+ "preferences": {
+ "precompile_float64": false
+ }
},
"Serialization": {
"uuid": "9e88b42a-f829-5b0c-bbe9-9e923198166b",
diff --git a/pysr/mutations.py b/pysr/mutations.py
new file mode 100644
index 000000000..57cc5052c
--- /dev/null
+++ b/pysr/mutations.py
@@ -0,0 +1,187 @@
+"""Mutation configurations for :class:`PySRRegressor`."""
+
+from __future__ import annotations
+
+from abc import ABC, abstractmethod
+from collections.abc import Mapping, Sequence
+from dataclasses import dataclass
+
+from .julia_helpers import jl_array
+from .julia_import import AnyValue, SymbolicRegression, jl
+
+_LEGACY_MUTATION_PARAMETERS = (
+ "weight_add_node",
+ "weight_insert_node",
+ "weight_delete_node",
+ "weight_do_nothing",
+ "weight_mutate_constant",
+ "weight_mutate_operator",
+ "weight_mutate_feature",
+ "weight_swap_operands",
+ "weight_rotate_tree",
+ "weight_randomize",
+ "weight_simplify",
+ "weight_optimize",
+ "weight_backsolve",
+)
+
+
+def create_mutations(mutation_pairs: Sequence[tuple[AnyValue, float]]) -> AnyValue:
+ """Create a Julia vector of weighted mutations."""
+ pair_type = jl.Pair[SymbolicRegression.AbstractMutation, jl.Float64]
+ return jl_array(mutation_pairs, dtype=pair_type)
+
+
+def convert_mutations(
+ mutation_weights: Mapping[AbstractMutation, float],
+) -> AnyValue:
+ """Convert Python mutation weights to their Julia representation."""
+ return create_mutations(
+ [
+ jl.Pair(mutation.julia_mutation(), weight)
+ for mutation, weight in mutation_weights.items()
+ ]
+ )
+
+
+class AbstractMutation(ABC):
+ """Base class for mutation configurations."""
+
+ @abstractmethod
+ def julia_mutation(self) -> AnyValue:
+ """Create the corresponding SymbolicRegression.jl mutation."""
+ pass # pragma: no cover
+
+
+@dataclass(frozen=True)
+class _ParameterlessMutation(AbstractMutation):
+ def julia_mutation(self) -> AnyValue:
+ return getattr(SymbolicRegression, type(self).__name__)()
+
+
+@dataclass(frozen=True)
+class ConstantMutation(AbstractMutation):
+ """Perturb a constant.
+
+ SymbolicRegression.jl default weight: ``0.0346``.
+
+ Defaults match SymbolicRegression.jl.
+ """
+
+ perturbation_factor: float = 0.086
+ probability_negate: float = 0.01
+
+ def julia_mutation(self) -> AnyValue:
+ return SymbolicRegression.ConstantMutation(
+ perturbation_factor=self.perturbation_factor,
+ probability_negate=self.probability_negate,
+ )
+
+
+@dataclass(frozen=True)
+class OperatorMutation(_ParameterlessMutation):
+ """Replace an operator with another operator of the same arity.
+
+ SymbolicRegression.jl default weight: ``0.293``.
+ """
+
+
+@dataclass(frozen=True)
+class FeatureMutation(_ParameterlessMutation):
+ """Change the feature referenced by a variable node.
+
+ SymbolicRegression.jl default weight: ``0.1``.
+ """
+
+
+@dataclass(frozen=True)
+class SwapOperandsMutation(_ParameterlessMutation):
+ """Swap the operands of a binary operator.
+
+ SymbolicRegression.jl default weight: ``0.198``.
+ """
+
+
+@dataclass(frozen=True)
+class AddNodeMutation(_ParameterlessMutation):
+ """Append a node to the expression.
+
+ SymbolicRegression.jl default weight: ``2.47``.
+ """
+
+
+@dataclass(frozen=True)
+class InsertNodeMutation(_ParameterlessMutation):
+ """Insert a node above an existing node.
+
+ SymbolicRegression.jl default weight: ``0.0112``.
+ """
+
+
+@dataclass(frozen=True)
+class DeleteNodeMutation(_ParameterlessMutation):
+ """Delete a node from the expression.
+
+ SymbolicRegression.jl default weight: ``0.870``.
+ """
+
+
+@dataclass(frozen=True)
+class RotateTreeMutation(_ParameterlessMutation):
+ """Rotate a subtree.
+
+ SymbolicRegression.jl default weight: ``4.26``.
+ """
+
+
+@dataclass(frozen=True)
+class BacksolveMutation(AbstractMutation):
+ """Fit a replacement expression by backsolving through the expression.
+
+ SymbolicRegression.jl default weight: ``0.0``.
+ """
+
+ max_library_size: int = 500
+ max_terms: int = 8
+ min_improvement: float = 1e-3
+ node_attempts: int = 8
+
+ def julia_mutation(self) -> AnyValue:
+ return SymbolicRegression.BacksolveMutation(
+ max_library_size=self.max_library_size,
+ max_terms=self.max_terms,
+ min_improvement=self.min_improvement,
+ node_attempts=self.node_attempts,
+ )
+
+
+@dataclass(frozen=True)
+class SimplifyMutation(_ParameterlessMutation):
+ """Simplify constant parts of the expression.
+
+ SymbolicRegression.jl default weight: ``0.00209``.
+ """
+
+
+@dataclass(frozen=True)
+class RandomizeMutation(_ParameterlessMutation):
+ """Replace the expression with a random expression.
+
+ SymbolicRegression.jl default weight: ``0.000502``.
+ """
+
+
+@dataclass(frozen=True)
+class OptimizeMutation(_ParameterlessMutation):
+ """Optimize constants as a mutation.
+
+ SymbolicRegression.jl default weight: ``0.0``.
+ """
+
+
+@dataclass(frozen=True)
+class DoNothingMutation(_ParameterlessMutation):
+ """Leave the expression unchanged.
+
+ SymbolicRegression.jl default weight: ``0.273``.
+ """
diff --git a/pysr/param_groupings.yml b/pysr/param_groupings.yml
index f4b769c27..70cd1b93c 100644
--- a/pysr/param_groupings.yml
+++ b/pysr/param_groupings.yml
@@ -4,6 +4,7 @@
- unary_operators
- operators
- expression_spec
+ - type_spec
- maxsize
- maxdepth
- Setting the Search Size:
@@ -33,6 +34,10 @@
- adaptive_parsimony_scaling
- should_simplify
- Mutations:
+ - default_mutations
+ - mutations
+ - default_plugins
+ - plugins
- weight_add_node
- weight_insert_node
- weight_delete_node
@@ -45,6 +50,7 @@
- weight_randomize
- weight_simplify
- weight_optimize
+ - weight_backsolve
- crossover_probability
- annealing
- alpha
@@ -101,6 +107,8 @@
- progress
- logger_spec
- input_stream
+ - use_tracing
+ - tracing_file
- Environment:
- temp_equation_file
- tempdir
diff --git a/pysr/plugins.py b/pysr/plugins.py
new file mode 100644
index 000000000..ade737c97
--- /dev/null
+++ b/pysr/plugins.py
@@ -0,0 +1,89 @@
+"""Plugin configurations for :class:`PySRRegressor`."""
+
+from __future__ import annotations
+
+from abc import ABC, abstractmethod
+from dataclasses import dataclass
+from typing import Literal
+
+from .julia_import import AnyValue, SymbolicRegression, jl
+
+
+class AbstractPlugin(ABC):
+ """Base class for plugin configurations."""
+
+ @abstractmethod
+ def julia_plugin(self) -> AnyValue:
+ """Create the corresponding SymbolicRegression.jl plugin."""
+ pass # pragma: no cover
+
+
+@dataclass(frozen=True)
+class SimulatedAnnealingPlugin(AbstractPlugin):
+ """Apply simulated annealing during mutation acceptance.
+
+ Enabled by default through ``annealing=True``.
+ """
+
+ alpha: float = 0.1
+
+ def julia_plugin(self) -> AnyValue:
+ return SymbolicRegression.SimulatedAnnealingPlugin(alpha=self.alpha)
+
+
+@dataclass(frozen=True)
+class AdaptiveParsimonyPlugin(AbstractPlugin):
+ """Apply frequency-based parsimony during selection and mutation.
+
+ Enabled by default in PySR through ``use_frequency=True`` and
+ ``use_frequency_in_tournament=True``.
+ """
+
+ tournament: bool = True
+ mutation_acceptance: bool = True
+
+ def julia_plugin(self) -> AnyValue:
+ return SymbolicRegression.AdaptiveParsimonyPlugin(
+ tournament=self.tournament,
+ mutation_acceptance=self.mutation_acceptance,
+ )
+
+
+@dataclass(frozen=True)
+class AdaptiveMutationWeightsPlugin(AbstractPlugin):
+ """Adapt mutation weights from their observed success rates.
+
+ Enabled by default in PySR. Pass an explicit ``default_plugins`` that
+ omits this plugin to disable it.
+ Learned runtime state is reinitialized for each call to ``fit``.
+ """
+
+ smoothing: float = 0.02
+ floor: float = 0.05
+ reward: Literal["cost", "loss"] = "cost"
+
+ def julia_plugin(self) -> AnyValue:
+ return SymbolicRegression.AdaptiveMutationWeightsPlugin(
+ smoothing=self.smoothing,
+ floor=self.floor,
+ reward=jl.Symbol(self.reward),
+ )
+
+
+@dataclass(frozen=True)
+class MutationBurstPlugin(AbstractPlugin):
+ """Retry rejected mutations and optionally chain accepted mutations.
+
+ Disabled by default in PySR.
+ """
+
+ retry_attempts: int = 4
+ compound_probability: float = 0.25
+ compound_max_steps: int = 2
+
+ def julia_plugin(self) -> AnyValue:
+ return SymbolicRegression.MutationBurstPlugin(
+ retry_attempts=self.retry_attempts,
+ compound_probability=self.compound_probability,
+ compound_max_steps=self.compound_max_steps,
+ )
diff --git a/pysr/sr.py b/pysr/sr.py
index f09b8bc50..fa4b46d6f 100644
--- a/pysr/sr.py
+++ b/pysr/sr.py
@@ -10,8 +10,9 @@
import sys
import tempfile
import warnings
-from collections.abc import Callable
+from collections.abc import Callable, Mapping, Sequence
from dataclasses import dataclass, fields
+from functools import wraps
from io import StringIO
from multiprocessing import cpu_count
from pathlib import Path
@@ -38,10 +39,10 @@
from .expression_specs import (
AbstractExpressionSpec,
ExpressionSpec,
- ParametricExpressionSpec,
- parametric_expression_deprecation_warning,
+ TemplateExpressionSpec,
)
from .feature_selection import run_feature_selection
+from .interrupt import _external_stop_signal_context
from .julia_extensions import load_required_packages
from .julia_helpers import (
_escape_filename,
@@ -50,10 +51,33 @@
jl_deserialize,
jl_is_function,
jl_named_tuple,
+ jl_numpy_array,
jl_serialize,
)
from .julia_import import AnyValue, SymbolicRegression, VectorValue, jl
from .logger_specs import AbstractLoggerSpec
+from .mutations import (
+ _LEGACY_MUTATION_PARAMETERS,
+ AbstractMutation,
+ convert_mutations,
+)
+from .plugins import AbstractPlugin
+from .type_specs import (
+ TypeSpec,
+ _TypeSpecRuntime,
+ _TypeSpecRuntimeDefinition,
+ compile_type_spec_runtime_for_model,
+ create_type_spec_addprocs_function,
+ create_type_spec_exports,
+ load_type_spec_runtime,
+ object_array_1d,
+ prepare_type_spec_fit_data,
+ prepare_type_spec_prediction_data,
+ type_spec_to_julia_array,
+ validate_type_spec_model_configuration,
+ validate_type_spec_options,
+ validate_type_spec_runtime,
+)
from .utils import (
ArrayLike,
PathLike,
@@ -76,8 +100,11 @@
from typing_extensions import List
+_CHECKPOINT_SCHEMA_VERSION = 3
+
ALREADY_RAN = False
+
pysr_logger = logging.getLogger(__name__)
@@ -103,7 +130,7 @@ def _process_constraints(
"One typical constraint is to use `constraints={..., '^': (-1, 1)}`, which "
"will allow arbitrary-complexity base (-1) but only powers such as "
"a constant or variable (1). "
- "For more tips, please see https://ai.damtp.cam.ac.uk/pysr/tuning/"
+ "For more tips, please see https://pysr.ai/tuning"
)
# Create default constraint tuple with -1 for each argument
constraints[op] = tuple([-1] * arity)
@@ -141,7 +168,7 @@ def _process_constraints(
def _maybe_create_inline_operators(
operators: dict[int, list[str]],
extra_sympy_mappings: dict[str, Callable] | None,
- expression_spec: AbstractExpressionSpec,
+ supports_sympy: bool,
) -> dict[int, list[str]]:
operators = {arity: op_list.copy() for arity, op_list in operators.items()}
@@ -150,24 +177,22 @@ def _maybe_create_inline_operators(
is_user_defined_operator = "(" in op
if is_user_defined_operator:
- jl.seval(op)
- # Cut off from the first non-alphanumeric char:
- first_non_char = [j for j, char in enumerate(op) if char == "("][0]
- function_name = op[:first_non_char]
- # Assert that function_name only contains
- # alphabetical characters, numbers,
- # and underscores:
+ function = jl.seval(op)
+ if not jl_is_function(function):
+ raise ValueError(
+ "Custom operator definitions must evaluate to a Julia function."
+ )
+ function_name = str(jl.Base.nameof(function))
if not re.match(r"^[a-zA-Z0-9_]+$", function_name):
raise ValueError(
- f"Invalid function name {function_name}. "
- "Only alphanumeric characters, numbers, "
- "and underscores are allowed."
+ "Custom operators must define a named Julia function whose name "
+ "contains only alphanumeric characters, numbers, and underscores."
)
missing_sympy_mapping = (
extra_sympy_mappings is None
or function_name not in extra_sympy_mappings
)
- if missing_sympy_mapping and expression_spec.supports_sympy:
+ if missing_sympy_mapping and supports_sympy:
raise ValueError(
f"Custom function {function_name} is not defined in `extra_sympy_mappings`. "
"You can define it with, "
@@ -179,6 +204,41 @@ def _maybe_create_inline_operators(
return operators
+def _create_julia_operators_and_loss_functions(
+ operators: dict[int, list[str]],
+ extra_sympy_mappings: dict[str, Callable] | None,
+ supports_sympy: bool,
+ elementwise_loss: str | None,
+ loss_function: str | None,
+ loss_function_expression: str | None,
+) -> tuple[dict[int, list[str]], AnyValue | None, AnyValue | None, AnyValue | None]:
+ operators = _maybe_create_inline_operators(
+ operators=operators,
+ extra_sympy_mappings=extra_sympy_mappings,
+ supports_sympy=supports_sympy,
+ )
+
+ def eval_source(source: str) -> Any:
+ return jl.seval(source)
+
+ def eval_objective(source: str | None, knob: str) -> AnyValue | None:
+ if source is None:
+ return None
+ loss = eval_source(str(source))
+ if not jl_is_function(loss):
+ raise ValueError(f"`{knob}` must evaluate to a callable Julia function.")
+ return loss
+
+ custom_loss = (
+ None if elementwise_loss is None else eval_source(str(elementwise_loss))
+ )
+ custom_full_objective = eval_objective(loss_function, "loss_function")
+ custom_loss_expression = eval_objective(
+ loss_function_expression, "loss_function_expression"
+ )
+ return operators, custom_loss, custom_full_objective, custom_loss_expression
+
+
def _check_assertions(
X,
use_custom_variable_names,
@@ -188,6 +248,7 @@ def _check_assertions(
y,
X_units,
y_units,
+ supports_sympy,
):
# Check for potential errors before they happen
assert len(X.shape) == 2
@@ -197,7 +258,8 @@ def _check_assertions(
assert weights.shape == y.shape
assert X.shape[0] == weights.shape[0]
if use_custom_variable_names:
- assert len(variable_names) == X.shape[1]
+ if len(variable_names) != X.shape[1]:
+ raise ValueError("`variable_names` must contain one name per feature.")
# Check none of the variable names are function names:
for var_name in variable_names:
# Check if alphanumeric only:
@@ -207,7 +269,8 @@ def _check_assertions(
"Only alphanumeric characters, numbers, "
"and underscores are allowed."
)
- assert_valid_sympy_symbol(var_name)
+ if supports_sympy:
+ assert_valid_sympy_symbol(var_name)
if (
isinstance(complexity_of_variables, list)
and len(complexity_of_variables) != X.shape[1]
@@ -343,6 +406,41 @@ def _validate_export_mappings(extra_jax_mappings, extra_torch_mappings):
# Class validation constants
VALID_OPTIMIZER_ALGORITHMS = ["BFGS", "NelderMead"]
+_WARM_START_MUTABLE_STATE = (
+ "equations_",
+ "equation_file_contents_",
+ "julia_state_stream_",
+ "julia_options_stream_",
+ "selection_mask_",
+ "feature_names_in_",
+ "display_feature_names_in_",
+ "complexity_of_variables_",
+ "X_units_",
+ "y_units_",
+)
+
+
+def _rollback_failed_warm_start(
+ fit_method: Callable[..., Any],
+) -> Callable[..., Any]:
+ @wraps(fit_method)
+ def wrapped(self, *args, **kwargs):
+ previous_state = None
+ if self.warm_start and getattr(self, "julia_state_stream_", None) is not None:
+ previous_state = self.__dict__.copy()
+ for name in _WARM_START_MUTABLE_STATE:
+ if name in previous_state:
+ previous_state[name] = copy.deepcopy(previous_state[name])
+ try:
+ return fit_method(self, *args, **kwargs)
+ except BaseException:
+ if previous_state is not None:
+ self.__dict__.clear()
+ self.__dict__.update(previous_state)
+ raise
+
+ return wrapped
+
@dataclass
class _DynamicallySetParams:
@@ -369,7 +467,7 @@ class PySRRegressor(MultiOutputMixin, RegressorMixin, BaseEstimator):
Most default parameters have been tuned over several example equations,
but you should adjust `niterations`, `binary_operators`, `unary_operators`
to your requirements. You can view more detailed explanations of the options
- on the [options page](https://ai.damtp.cam.ac.uk/pysr/options) of the
+ on the [options page](https://pysr.ai/options) of the
documentation.
Parameters
@@ -389,7 +487,7 @@ class PySRRegressor(MultiOutputMixin, RegressorMixin, BaseEstimator):
most accurate model.
binary_operators : list[str]
List of strings for binary operators used in the search.
- See the [operators page](https://ai.damtp.cam.ac.uk/pysr/operators/)
+ See the [operators page](https://pysr.ai/operators)
for more details.
Default is `["+", "-", "*", "/"]`.
unary_operators : list[str]
@@ -408,6 +506,10 @@ class PySRRegressor(MultiOutputMixin, RegressorMixin, BaseEstimator):
`TemplateExpressionSpec(...)` which allows you to specify
a custom template for the expressions.
Default is `ExpressionSpec()`.
+ type_spec : TypeSpec
+ Declarative custom value type. PySR builds an isolated Julia wrapper,
+ interface methods, operators, and loss from this specification.
+ Default is `None`, which uses the numeric path.
niterations : int
Number of iterations of the algorithm to run. The best
equations are printed and migrate between populations at the
@@ -567,7 +669,7 @@ class PySRRegressor(MultiOutputMixin, RegressorMixin, BaseEstimator):
(requires `annealing` to be `True`).
Default is `3.17`.
annealing : bool
- Whether to use annealing. Default is `False`.
+ Whether to use annealing. Default is `True`.
early_stop_condition : float | str
Stop the search early if this loss is reached. You may also
pass a string containing a Julia function which
@@ -588,50 +690,61 @@ class PySRRegressor(MultiOutputMixin, RegressorMixin, BaseEstimator):
fraction_replaced_guesses : float
How much of the population to replace with migrating equations from
guesses. Default is `0.001`.
- weight_add_node : float
+ weight_add_node : float | None
Relative likelihood for mutation to add a node.
- Default is `2.47`.
- weight_insert_node : float
+ Default is `None` (mapping to `2.47`).
+ weight_insert_node : float | None
Relative likelihood for mutation to insert a node.
- Default is `0.0112`.
- weight_delete_node : float
+ Default is `None` (mapping to `0.0112`).
+ weight_delete_node : float | None
Relative likelihood for mutation to delete a node.
- Default is `0.870`.
- weight_do_nothing : float
+ Default is `None` (mapping to `0.870`).
+ weight_do_nothing : float | None
Relative likelihood for mutation to leave the individual.
- Default is `0.273`.
- weight_mutate_constant : float
+ Default is `None` (mapping to `0.273`).
+ weight_mutate_constant : float | None
Relative likelihood for mutation to change the constant slightly
in a random direction.
- Default is `0.0346`.
- weight_mutate_operator : float
+ Default is `None` (mapping to `0.0346`).
+ weight_mutate_operator : float | None
Relative likelihood for mutation to swap an operator.
- Default is `0.293`.
- weight_mutate_feature : float
+ Default is `None` (mapping to `0.293`).
+ weight_mutate_feature : float | None
Relative likelihood for mutation to change which feature a variable node references.
- Default is `0.1`.
- weight_swap_operands : float
+ Default is `None` (mapping to `0.1`).
+ weight_swap_operands : float | None
Relative likehood for swapping operands in binary operators.
- Default is `0.198`.
- weight_rotate_tree : float
+ Default is `None` (mapping to `0.198`).
+ weight_rotate_tree : float | None
How often to perform a tree rotation at a random node.
- Default is `4.26`.
- weight_randomize : float
+ Default is `None` (mapping to `4.26`).
+ weight_randomize : float | None
Relative likelihood for mutation to completely delete and then
randomly generate the equation
- Default is `0.000502`.
- weight_simplify : float
+ Default is `None` (mapping to `0.000502`).
+ weight_simplify : float | None
Relative likelihood for mutation to simplify constant parts by evaluation
- Default is `0.00209`.
- weight_optimize: float
+ Default is `None` (mapping to `0.00209`).
+ weight_optimize: float | None
Constant optimization can also be performed as a mutation, in addition to
the normal strategy controlled by `optimize_probability` which happens
every iteration. Using it as a mutation is useful if you want to use
a large `ncycles_periteration`, and may not optimize very often.
- Default is `0.0`.
+ Default is `None` (mapping to `0.0`).
+ weight_backsolve : float | None
+ Relative likelihood for backsolve mutation. To configure its parameters,
+ pass `BacksolveMutation(...)` through `mutations` instead.
+ Default is `None` (mapping to `0.0`).
+ mutations : Mapping[AbstractMutation, float] | None
+ Mutation configurations and their weights. Entries override or extend
+ `default_mutations` by mutation type. Default is `None`.
+ default_mutations : Mapping[AbstractMutation, float] | None
+ Default mutation configurations and their weights. When provided, these
+ replace the SymbolicRegression.jl defaults and plugin-contributed mutations.
+ Default is `None`.
crossover_probability : float
Absolute probability of crossover-type genetic operation, instead of a mutation.
- Default is `0.0259`.
+ Default is `0.2`.
skip_mutation_failures : bool
Whether to skip mutation and crossover failures, rather than
simply re-sampling the current member.
@@ -690,11 +803,10 @@ class PySRRegressor(MultiOutputMixin, RegressorMixin, BaseEstimator):
Number of processes to use for parallelism. If `None`, defaults to `cpu_count()`.
Default is `None`.
cluster_manager : str
- For distributed computing, this sets the job queue system. Set
- to one of "slurm", "pbs", "lsf", "sge", "qrsh", "scyld", or
- "htc". If set to one of these, PySR will run in distributed
- mode, and use `procs` to figure out how many processes to launch.
- Default is `None`.
+ For distributed computing, this sets the job queue system. Set to
+ "slurm" to use an existing Slurm allocation, with `procs` equal to
+ its task count. Other supported values are "pbs", "lsf", "sge",
+ "qrsh", "scyld", and "htc". Default is `None`.
heap_size_hint_in_bytes : int
For multiprocessing, this sets the `--heap-size-hint` parameter
for new Julia processes. This can be configured when using
@@ -706,13 +818,17 @@ class PySRRegressor(MultiOutputMixin, RegressorMixin, BaseEstimator):
If a worker does not respond within this time, it will be restarted.
Default is `None`.
worker_imports : list[str] | None
- List of module names as strings to import in worker processes.
- For example, `["MyPackage", "OtherPackage"]` will run `using MyPackage, OtherPackage`
- in each worker process. Default is `None`.
+ Module names to import on multiprocessing workers. For example,
+ ``["MyPackage", "OtherPackage"]`` runs
+ ``using MyPackage, OtherPackage`` on every worker. TypeSpec preambles,
+ operators, objectives, and expression specifications may depend on
+ packages listed here; arbitrary bindings from ``Main`` are unavailable.
+ Default is ``None``.
batching : bool | "auto"
Whether to compare population members on small batches during
evolution. Still uses full dataset for comparing against hall
- of fame. "auto" enables batching for N>1000. Default is `"auto"`.
+ of fame. `"auto"` lets SymbolicRegression.jl choose based on the
+ dataset. Default is `"auto"`.
batch_size : int | None
The batch size to use if batching. If None, uses the
full dataset when N<=1000, 128 for N<5000, 256 for N<50000,
@@ -756,12 +872,17 @@ class PySRRegressor(MultiOutputMixin, RegressorMixin, BaseEstimator):
warm_start : bool
Tells fit to continue from where the last call to fit finished.
If false, each call to fit will be fresh, overwriting previous results.
+ Plugin runtime state is reinitialized for each call to `fit`.
Default is `False`.
- guesses : list[str] | list[list[str]] | list[dict[str, str]] | list[list[dict[str, str]]] | None
+ guesses : list[str] | list[list[str]] | list[dict[str, str | ArrayLike]] | list[list[dict[str, str | ArrayLike]]] | None
Initial guesses for expressions to seed the search. Examples:
`["x0 + x1", "x0^2"]`, `[["x0"], ["x1"]]` (multi-output),
- `[{"f": "#1 + #2"}]` (TemplateExpressionSpec where `#1`, `#2` are
- placeholders for the 1st, 2nd arguments of expression `f`).
+ or `[{"f": "#1 + #2", "p": [5.0, 10.0]}]` for a
+ `TemplateExpressionSpec` with component `f` and parameter vector `p`.
+ Here `#1` and `#2` are the first and second arguments of `f`.
+ Parameter values use the data precision or the model's `TypeSpec`, and
+ supplied vectors must match the lengths declared by the template.
+ Parameter names may be omitted to retain their normal initialization.
Default is `None`.
verbosity : int
What verbosity level to use. 0 means minimal print statements.
@@ -779,12 +900,30 @@ class PySRRegressor(MultiOutputMixin, RegressorMixin, BaseEstimator):
Logger specification for the Julia backend. See, for example,
`TensorBoardLoggerSpec`.
Default is `None`.
- input_stream : str
- The stream to read user input from. By default, this is `"stdin"`.
- If you encounter issues with reading from `stdin`, like a hang,
- you can simply pass `"devnull"` to this argument. You can also
- reference an arbitrary Julia object in the `Main` namespace.
- Default is `"stdin"`.
+ plugins : Sequence[AbstractPlugin] | None
+ Plugin configurations. Entries override or extend `default_plugins` by
+ plugin type. Default is `None`.
+ default_plugins : Sequence[AbstractPlugin] | None
+ Default plugin configurations. Default is `None`.
+ input_stream : str | None
+ The stream to read user input from, used for the `'q'` + ``
+ command that stops a search early. `None` reads from `"stdin"` when
+ attached to an interactive terminal and otherwise disables stdin
+ watching via `"devnull"` (for example in Jupyter, where typed input
+ never reaches the search). You can also pass `"stdin"` or `"devnull"`
+ explicitly, or reference an arbitrary Julia object in the `Main`
+ namespace. Default is `None`.
+ use_tracing : bool
+ Whether to write a JSONL trace of the search to `tracing_file`.
+ One record per line holds the live members of one population at one
+ iteration, along with every mutation, crossover, tuning and death
+ event. Uses JSON.jl, which is installed into the Julia environment
+ on first use.
+ Default is `False`.
+ tracing_file : str | Path
+ Where to write the trace when `use_tracing` is set. Any existing
+ file is overwritten when the search starts.
+ Default is `"pysr_trace.jsonl"`.
run_id : str
A unique identifier for the run. Will be generated using the
current date and time if not provided.
@@ -944,6 +1083,7 @@ class PySRRegressor(MultiOutputMixin, RegressorMixin, BaseEstimator):
logger_: AnyValue | None
equation_file_contents_: list[pd.DataFrame] | None
show_pickle_warnings_: bool
+ _type_spec_runtime_definition_: _TypeSpecRuntimeDefinition
def __init__(
self,
@@ -953,6 +1093,7 @@ def __init__(
unary_operators: list[str] | None = None,
operators: dict[int, list[str]] | None = None,
expression_spec: AbstractExpressionSpec | None = None,
+ type_spec: TypeSpec | None = None,
niterations: int = 100,
populations: int = 31,
population_size: int = 27,
@@ -978,25 +1119,28 @@ def __init__(
use_frequency_in_tournament: bool = True,
adaptive_parsimony_scaling: float = 1040.0,
alpha: float = 3.17,
- annealing: bool = False,
+ annealing: bool = True,
early_stop_condition: float | str | None = None,
ncycles_per_iteration: int = 380,
fraction_replaced: float = 0.00036,
fraction_replaced_hof: float = 0.0614,
fraction_replaced_guesses: float = 0.001,
- weight_add_node: float = 2.47,
- weight_insert_node: float = 0.0112,
- weight_delete_node: float = 0.870,
- weight_do_nothing: float = 0.273,
- weight_mutate_constant: float = 0.0346,
- weight_mutate_operator: float = 0.293,
- weight_mutate_feature: float = 0.1,
- weight_swap_operands: float = 0.198,
- weight_rotate_tree: float = 4.26,
- weight_randomize: float = 0.000502,
- weight_simplify: float = 0.00209,
- weight_optimize: float = 0.0,
- crossover_probability: float = 0.0259,
+ weight_add_node: float | None = None,
+ weight_insert_node: float | None = None,
+ weight_delete_node: float | None = None,
+ weight_do_nothing: float | None = None,
+ weight_mutate_constant: float | None = None,
+ weight_mutate_operator: float | None = None,
+ weight_mutate_feature: float | None = None,
+ weight_swap_operands: float | None = None,
+ weight_rotate_tree: float | None = None,
+ weight_randomize: float | None = None,
+ weight_simplify: float | None = None,
+ weight_optimize: float | None = None,
+ weight_backsolve: float | None = None,
+ mutations: Mapping[AbstractMutation, float] | None = None,
+ default_mutations: Mapping[AbstractMutation, float] | None = None,
+ crossover_probability: float = 0.2,
skip_mutation_failures: bool = True,
migration: bool = True,
hof_migration: bool = True,
@@ -1035,8 +1179,10 @@ def __init__(
guesses: (
list[str]
| list[list[str]]
- | list[dict[str, str]]
- | list[list[dict[str, str]]]
+ | list[dict[str, str | ArrayLike]]
+ | list[list[dict[str, str | ArrayLike]]]
+ | list[AnyValue]
+ | list[list[AnyValue]]
| None
) = None,
verbosity: int = 1,
@@ -1044,7 +1190,11 @@ def __init__(
print_precision: int = 5,
progress: bool = True,
logger_spec: AbstractLoggerSpec | None = None,
- input_stream: str = "stdin",
+ plugins: Sequence[AbstractPlugin] | None = None,
+ default_plugins: Sequence[AbstractPlugin] | None = None,
+ input_stream: str | None = None,
+ use_tracing: bool = False,
+ tracing_file: str | Path = "pysr_trace.jsonl",
run_id: str | None = None,
output_directory: str | None = None,
temp_equation_file: bool = False,
@@ -1067,6 +1217,7 @@ def __init__(
self.unary_operators = unary_operators
self.operators = operators
self.expression_spec = expression_spec
+ self.type_spec = type_spec
self.niterations = niterations
self.populations = populations
self.population_size = population_size
@@ -1113,6 +1264,9 @@ def __init__(
self.weight_randomize = weight_randomize
self.weight_simplify = weight_simplify
self.weight_optimize = weight_optimize
+ self.weight_backsolve = weight_backsolve
+ self.mutations = mutations
+ self.default_mutations = default_mutations
self.crossover_probability = crossover_probability
self.skip_mutation_failures = skip_mutation_failures
# -- Migration parameters
@@ -1159,7 +1313,11 @@ def __init__(
self.print_precision = print_precision
self.progress = progress
self.logger_spec = logger_spec
+ self.plugins = plugins
+ self.default_plugins = default_plugins
self.input_stream = input_stream
+ self.use_tracing = use_tracing
+ self.tracing_file = tracing_file
# - Project management
self.run_id = run_id
self.output_directory = output_directory
@@ -1202,7 +1360,7 @@ def __init__(
elif k == "julia_project":
warnings.warn(
"The `julia_project` parameter has been deprecated. To use a custom "
- "julia project, please see `https://ai.damtp.cam.ac.uk/pysr/backend`.",
+ "julia project, please see `https://pysr.ai/backend`.",
FutureWarning,
)
elif k == "julia_kwargs":
@@ -1301,11 +1459,19 @@ def from_file(
if "equations_" not in model.__dict__ or model.equations_ is None:
model.refresh()
+ else:
+ model._restore_julia_backed_columns()
- if model.expression_spec is not None:
+ if (
+ not isinstance(
+ model.expression_spec_, (ExpressionSpec, TemplateExpressionSpec)
+ )
+ and not model._has_fitted_type_spec()
+ ):
warnings.warn(
- "Loading model from checkpoint file with a non-default expression spec "
- "is not fully supported as it relies on dynamic objects. This may result in unexpected behavior.",
+ "Loading a checkpoint with a custom expression spec is not fully "
+ "supported, as it relies on dynamic Julia objects. This may result "
+ "in unexpected behavior.",
)
return model
@@ -1314,6 +1480,11 @@ def from_file(
f"Checkpoint file {pkl_filename} does not exist. "
"Attempting to recreate model from scratch..."
)
+ if pysr_kwargs.get("type_spec") is not None:
+ raise ValueError(
+ "TypeSpec models require the original `checkpoint.pkl`; "
+ "CSV-only reconstruction is not supported."
+ )
csv_filename = Path(run_directory) / "hall_of_fame.csv"
csv_filename_bak = Path(run_directory) / "hall_of_fame.csv.bak"
if not csv_filename.exists() and not csv_filename_bak.exists():
@@ -1436,27 +1607,33 @@ def __getstate__(self) -> dict[str, Any]:
warnings.warn(warn_msg)
else:
pysr_logger.debug(warn_msg)
- state_keys_to_clear = state_keys_containing_lambdas
- state_keys_to_clear.append("logger_")
+ state_keys_to_clear = [*state_keys_containing_lambdas, "logger_"]
pickled_state = {
key: (None if key in state_keys_to_clear else value)
for key, value in state.items()
}
+ pickled_state["_checkpoint_schema_version"] = _CHECKPOINT_SCHEMA_VERSION
if ("equations_" in pickled_state) and (
pickled_state["equations_"] is not None
):
pickled_state["output_torch_format"] = False
pickled_state["output_jax_format"] = False
+ unpicklable_columns = ["jax_format", "torch_format"]
+ if self._has_julia_backed_equations():
+ # Live Julia objects cannot be unpickled in a fresh process
+ # before their Julia definitions exist; these columns are
+ # rebuilt from `julia_state_` via `refresh()`.
+ unpicklable_columns += ["julia_expression", "lambda_format"]
if self.nout_ == 1:
pickled_columns = ~pickled_state["equations_"].columns.isin(
- ["jax_format", "torch_format"]
+ unpicklable_columns
)
pickled_state["equations_"] = (
pickled_state["equations_"].loc[:, pickled_columns].copy()
)
else:
pickled_columns = [
- ~dataframe.columns.isin(["jax_format", "torch_format"])
+ ~dataframe.columns.isin(unpicklable_columns)
for dataframe in pickled_state["equations_"]
]
pickled_state["equations_"] = [
@@ -1467,19 +1644,41 @@ def __getstate__(self) -> dict[str, Any]:
]
return pickled_state
+ def __setstate__(self, state: dict[str, Any]) -> None:
+ schema_version = state.pop("_checkpoint_schema_version", None)
+ state.setdefault("type_spec", None)
+ if schema_version != _CHECKPOINT_SCHEMA_VERSION:
+ raise ValueError(
+ "Unsupported PySR checkpoint schema: "
+ f"expected {_CHECKPOINT_SCHEMA_VERSION}, found {schema_version!r}."
+ )
+ self.__dict__.update(state)
+
def _checkpoint(self):
"""Save the model's current state to a checkpoint file.
This should only be used internally by PySRRegressor.
"""
- # Save model state:
+ checkpoint_path = self.get_pkl_filename()
+ previous_show_pickle_warnings = getattr(self, "show_pickle_warnings_", True)
+ # Same directory as the destination, so `os.replace` stays atomic and a
+ # failed pickle cannot truncate an existing checkpoint:
+ temporary_path = checkpoint_path.with_name(checkpoint_path.name + ".tmp")
self.show_pickle_warnings_ = False
- with open(self.get_pkl_filename(), "wb") as f:
+ try:
+ with open(temporary_path, "wb") as checkpoint_file:
+ try:
+ pkl.dump(self, checkpoint_file)
+ except Exception as e:
+ pysr_logger.debug(f"Error checkpointing model: {e}")
+ return
+ os.replace(temporary_path, checkpoint_path)
+ finally:
+ self.show_pickle_warnings_ = previous_show_pickle_warnings
try:
- pkl.dump(self, f)
+ temporary_path.unlink(missing_ok=True)
except Exception as e:
- pysr_logger.debug(f"Error checkpointing model: {e}")
- self.show_pickle_warnings_ = True
+ pysr_logger.debug(f"Error cleaning up temporary checkpoint file: {e}")
def get_pkl_filename(self) -> Path:
path = Path(self.output_directory_) / self.run_id_ / "checkpoint.pkl"
@@ -1498,14 +1697,22 @@ def equations(self): # pragma: no cover
@property
def julia_options_(self):
"""The deserialized julia options."""
- return jl_deserialize(self.julia_options_stream_)
+ stream = getattr(self, "julia_options_stream_", None)
+ if stream is None:
+ return None
+ self._define_julia_expression_types()
+ return jl_deserialize(stream)
@property
def julia_state_(self):
"""The deserialized state."""
+ stream = getattr(self, "julia_state_stream_", None)
+ if stream is None:
+ return None
+ self._define_julia_expression_types()
return cast(
Union[Tuple[VectorValue, AnyValue], None],
- jl_deserialize(self.julia_state_stream_),
+ jl_deserialize(stream),
)
@property
@@ -1522,6 +1729,39 @@ def raw_julia_state_(self):
def expression_spec_(self):
return self.expression_spec or ExpressionSpec()
+ def _has_fitted_type_spec(self) -> bool:
+ return hasattr(self, "_type_spec_runtime_definition_")
+
+ def _supports_export(self, format: str) -> bool:
+ return not self._has_fitted_type_spec() and bool(
+ getattr(self.expression_spec_, f"supports_{format}")
+ )
+
+ def _has_julia_backed_equations(self) -> bool:
+ """Whether `equations_` holds live Julia objects that cannot be pickled."""
+ return self._has_fitted_type_spec() or isinstance(
+ self.expression_spec_, TemplateExpressionSpec
+ )
+
+ def _define_julia_expression_types(self) -> None:
+ """Define the Julia types needed to deserialize `julia_state_stream_`."""
+ if self._has_fitted_type_spec():
+ self._load_type_spec_runtime()
+ elif isinstance(self.expression_spec_, TemplateExpressionSpec):
+ self.expression_spec_.julia_expression_spec()
+
+ def _restore_julia_backed_columns(self) -> None:
+ """Rebuild the Julia-backed equation columns dropped by pickling."""
+ if not self._has_julia_backed_equations():
+ return
+ equations = getattr(self, "equations_", None)
+ frames = equations if isinstance(equations, list) else [equations]
+ if any(
+ isinstance(frame, pd.DataFrame) and "lambda_format" not in frame.columns
+ for frame in frames
+ ):
+ self.refresh()
+
def get_best(
self, index: int | list[int] | None = None
) -> pd.Series | list[pd.Series]:
@@ -1547,6 +1787,7 @@ def get_best(
Raised when an invalid model selection strategy is provided.
"""
check_is_fitted(self, attributes=["equations_"])
+ self._restore_julia_backed_columns()
if index is not None:
if isinstance(self.equations_, list):
@@ -1611,6 +1852,63 @@ def _setup_equation_file(self):
def _clear_equation_file_contents(self):
self.equation_file_contents_ = None
+ def _operators_from_params(self) -> dict[int, list[str]]:
+ if self.operators is not None:
+ return {arity: values.copy() for arity, values in self.operators.items()}
+ operators = {
+ 2: (
+ self.binary_operators.copy()
+ if self.binary_operators is not None
+ else ["+", "-", "/", "*"]
+ )
+ }
+ if self.unary_operators is not None:
+ operators[1] = self.unary_operators.copy()
+ return operators
+
+ def _compile_type_spec_runtime(
+ self, operators: dict[int, list[str]]
+ ) -> _TypeSpecRuntimeDefinition:
+ return compile_type_spec_runtime_for_model(self, operators)
+
+ def _load_type_spec_runtime(
+ self,
+ *,
+ for_fit: bool = False,
+ operators: dict[int, list[str]] | None = None,
+ ) -> _TypeSpecRuntime:
+ if not for_fit:
+ check_is_fitted(self, attributes=["_type_spec_runtime_definition_"])
+ return load_type_spec_runtime(self._type_spec_runtime_definition_)
+
+ assert operators is not None
+ definition = self._compile_type_spec_runtime(operators)
+ fitted_definition = getattr(self, "_type_spec_runtime_definition_", None)
+ if (
+ self.warm_start
+ and fitted_definition is not None
+ and definition.fingerprint != fitted_definition.fingerprint
+ ):
+ raise ValueError(
+ "Cannot warm-start after changing the TypeSpec configuration. "
+ "Start a new search with `warm_start=False`."
+ )
+ runtime = load_type_spec_runtime(definition)
+ if (
+ fitted_definition is None
+ or definition.fingerprint != fitted_definition.fingerprint
+ ):
+ validate_type_spec_runtime(runtime)
+ return runtime
+
+ def _julia_expression_spec(
+ self,
+ type_spec_runtime: _TypeSpecRuntime | None = None,
+ ) -> AnyValue:
+ if type_spec_runtime is not None:
+ return type_spec_runtime.expression_spec
+ return self.expression_spec_.julia_expression_spec()
+
def _validate_and_modify_params(self) -> _DynamicallySetParams:
"""
Ensure parameters passed at initialization are valid.
@@ -1625,8 +1923,29 @@ def _validate_and_modify_params(self) -> _DynamicallySetParams:
values. For example, default parameters are set here
when a parameter is left set to `None`.
"""
- # Immutable parameter validation
- # Ensure instance parameters are allowable values:
+ if (
+ self.warm_start
+ and getattr(self, "julia_state_stream_", None) is not None
+ and (self.type_spec is not None) != self._has_fitted_type_spec()
+ ):
+ raise ValueError(
+ "Cannot warm-start after enabling or disabling TypeSpec. "
+ "Start a new search with `warm_start=False`."
+ )
+
+ if self.type_spec is not None:
+ validate_type_spec_model_configuration(self)
+ legacy_mutation_weights_used = any(
+ getattr(self, parameter) is not None
+ for parameter in _LEGACY_MUTATION_PARAMETERS
+ )
+ if legacy_mutation_weights_used and (
+ self.default_mutations is not None or self.mutations is not None
+ ):
+ raise ValueError(
+ "Cannot combine legacy `weight_*` parameters with "
+ "`default_mutations` or `mutations`."
+ )
# Validate operators vs binary_operators/unary_operators mutual exclusion
if self.operators is not None:
@@ -1660,7 +1979,7 @@ def _validate_and_modify_params(self) -> _DynamicallySetParams:
)
param_container = _DynamicallySetParams(
- operators={2: ["+", "*", "-", "/"]},
+ operators={2: ["+", "-", "/", "*"]},
maxdepth=self.maxsize,
constraints={},
batch_size=None,
@@ -1670,17 +1989,7 @@ def _validate_and_modify_params(self) -> _DynamicallySetParams:
)
# Convert binary_operators/unary_operators to operators format if needed
- if self.operators is None:
- # Build operators dict from binary_operators and unary_operators
- operators_dict = {}
- if self.binary_operators is not None:
- operators_dict[2] = self.binary_operators.copy()
- else:
- # Keep default binary operators
- operators_dict[2] = ["+", "*", "-", "/"]
- if self.unary_operators is not None:
- operators_dict[1] = self.unary_operators.copy()
- param_container.operators = operators_dict
+ param_container.operators = self._operators_from_params()
for param_name in map(lambda x: x.name, fields(_DynamicallySetParams)):
user_param_value = getattr(self, param_name)
@@ -1765,6 +2074,31 @@ def _validate_and_set_fit_params(
Validated units for `y`.
"""
+ if (
+ complexity_of_variables is not None
+ and self.complexity_of_variables is not None
+ ):
+ raise ValueError(
+ "You cannot set `complexity_of_variables` at both `fit` and `__init__`. "
+ "Pass it at `__init__` to set it to global default, OR use `fit` to set it for "
+ "each variable individually."
+ )
+ elif complexity_of_variables is None:
+ complexity_of_variables = self.complexity_of_variables
+
+ if self.type_spec is not None:
+ return prepare_type_spec_fit_data(
+ self,
+ X,
+ y,
+ Xresampled,
+ weights,
+ variable_names,
+ complexity_of_variables,
+ X_units,
+ y_units,
+ )
+
if isinstance(X, pd.DataFrame):
if variable_names:
variable_names = None
@@ -1789,22 +2123,6 @@ def _validate_and_set_fit_params(
"Please use valid names instead."
)
- if (
- complexity_of_variables is not None
- and self.complexity_of_variables is not None
- ):
- raise ValueError(
- "You cannot set `complexity_of_variables` at both `fit` and `__init__`. "
- "Pass it at `__init__` to set it to global default, OR use `fit` to set it for "
- "each variable individually."
- )
- elif complexity_of_variables is not None:
- complexity_of_variables = complexity_of_variables
- elif self.complexity_of_variables is not None:
- complexity_of_variables = self.complexity_of_variables
- else:
- complexity_of_variables = None
-
# Data validation and feature name fetching via sklearn
# This method sets the n_features_in_ attribute
if Xresampled is not None:
@@ -1993,8 +2311,10 @@ def _run(
y: ndarray,
runtime_params: _DynamicallySetParams,
weights: ndarray | None,
- category: ndarray | None,
seed: int,
+ type_spec_runtime: _TypeSpecRuntime | None,
+ parallelism: Literal["serial", "multithreading", "multiprocessing"],
+ numprocs: int | None,
):
"""
Run the symbolic regression fitting process on the julia backend.
@@ -2012,12 +2332,14 @@ def _run(
Weight array of the same shape as `y`.
Each element is how to weight the mean-square-error loss
for that particular element of y.
- category : ndarray | None
- If `expression_spec` is a `ParametricExpressionSpec`, then this
- argument should be a list of integers representing the category
- of each sample in `X`.
seed : int
Random seed for julia backend process.
+ type_spec_runtime : _TypeSpecRuntime | None
+ Loaded TypeSpec runtime, or `None` for numeric searches.
+ parallelism : {"serial", "multithreading", "multiprocessing"}
+ Effective Julia backend execution mode.
+ numprocs : int | None
+ Effective number of Julia worker processes.
Returns
-------
@@ -2033,9 +2355,35 @@ def _run(
# every new instance of PySRRegressor
global ALREADY_RAN
- # These are the parameters which may be modified from the ones
- # specified in init, so we define them here locally:
- operators = runtime_params.operators
+ definition_module = (
+ jl.Main
+ if type_spec_runtime is None
+ else type_spec_runtime.configuration_module
+ )
+ type_spec_operator_functions: dict[int, tuple[AnyValue, ...]] | None = None
+ if type_spec_runtime is None:
+ supports_sympy = self.expression_spec_.supports_sympy
+ operators, custom_loss, custom_full_objective, custom_loss_expression = (
+ _create_julia_operators_and_loss_functions(
+ operators=runtime_params.operators,
+ extra_sympy_mappings=self.extra_sympy_mappings,
+ supports_sympy=supports_sympy,
+ elementwise_loss=self.elementwise_loss,
+ loss_function=self.loss_function,
+ loss_function_expression=self.loss_function_expression,
+ )
+ )
+ else:
+ type_spec_operator_functions = type_spec_runtime.operator_functions
+ operators = {
+ arity: list(names)
+ for arity, names in type_spec_runtime.operator_names.items()
+ }
+ custom_loss = type_spec_runtime.elementwise_loss
+ custom_full_objective = type_spec_runtime.loss_function
+ custom_loss_expression = type_spec_runtime.loss_function_expression
+ value_type = None if type_spec_runtime is None else type_spec_runtime.value_type
+ loss_type = None
constraints = runtime_params.constraints
nested_constraints = self.nested_constraints
@@ -2047,10 +2395,6 @@ def _run(
if not ALREADY_RAN and runtime_params.update_verbosity != 0:
pysr_logger.info("Compiling Julia backend...")
- parallelism, numprocs = _map_parallelism_params(
- self.parallelism, self.procs, getattr(self, "multithreading", None)
- )
-
if self.deterministic and parallelism != "serial":
raise ValueError(
"To ensure deterministic searches, you must set `parallelism='serial'`. "
@@ -2070,12 +2414,6 @@ def _run(
"To use cluster managers, you must set `parallelism='multiprocessing'`."
)
- # TODO(mcranmer): These functions should be part of this class.
- operators = _maybe_create_inline_operators(
- operators=operators,
- extra_sympy_mappings=self.extra_sympy_mappings,
- expression_spec=self.expression_spec_,
- )
if constraints is not None:
_constraints = _process_constraints(
operators=operators,
@@ -2095,7 +2433,7 @@ def _run(
max_arity = max(operators.keys()) if operators else 2
constraints_by_arity = {arity: None for arity in range(1, max_arity + 1)}
- # Parse dict into Julia Dict for nested constraints::
+ # Parse dict into Julia Dict for nested constraints:
if nested_constraints is not None:
nested_constraints_str = "Dict("
for outer_k, outer_v in nested_constraints.items():
@@ -2104,7 +2442,9 @@ def _run(
nested_constraints_str += f"({inner_k}) => {inner_v}, "
nested_constraints_str += "), "
nested_constraints_str += ")"
- nested_constraints = jl.seval(nested_constraints_str)
+ nested_constraints = jl.Base.include_string(
+ definition_module, nested_constraints_str, "PySR nested_constraints"
+ )
# Parse dict into Julia Dict for complexities:
if complexity_of_operators is not None:
@@ -2112,47 +2452,48 @@ def _run(
for k, v in complexity_of_operators.items():
complexity_of_operators_str += f"({k}) => {v}, "
complexity_of_operators_str += ")"
- complexity_of_operators = jl.seval(complexity_of_operators_str)
+ complexity_of_operators = jl.Base.include_string(
+ definition_module,
+ complexity_of_operators_str,
+ "PySR complexity_of_operators",
+ )
# TODO: Refactor this into helper function
if isinstance(complexity_of_variables, list):
complexity_of_variables = jl_array(complexity_of_variables)
- np_dtype = self._get_precision_mapped_dtype(np.array(X))
-
- custom_loss = jl.seval(
- str(self.elementwise_loss)
- if self.elementwise_loss is not None
- else "nothing"
+ np_dtype = (
+ None
+ if type_spec_runtime is not None
+ else self._get_precision_mapped_dtype(np.array(X))
)
+
if self.elementwise_loss is not None:
- _validate_elementwise_loss(
- custom_loss,
- has_weights=weights is not None,
- probe_value=np_dtype(1.0),
- )
+ if type_spec_runtime is None:
+ assert np_dtype is not None
+ _validate_elementwise_loss(
+ custom_loss,
+ has_weights=weights is not None,
+ probe_value=np_dtype(1.0),
+ )
- custom_full_objective = jl.seval(
- str(self.loss_function) if self.loss_function is not None else "nothing"
- )
if self.loss_function is not None:
_validate_custom_full_objective(custom_full_objective)
- custom_loss_expression = jl.seval(
- str(self.loss_function_expression)
- if self.loss_function_expression is not None
- else "nothing"
- )
if self.loss_function_expression is not None:
_validate_custom_expression_objective(custom_loss_expression)
- early_stop_condition = jl.seval(
- str(self.early_stop_condition)
- if self.early_stop_condition is not None
- else "nothing"
+ early_stop_condition = (
+ jl.seval(
+ str(self.early_stop_condition)
+ if self.early_stop_condition is not None
+ else "nothing"
+ )
+ if type_spec_runtime is None
+ else type_spec_runtime.early_stop_condition
)
- input_stream = jl.seval(self.input_stream)
+ input_stream = jl.seval(_resolve_input_stream(self.input_stream))
load_required_packages(
turbo=self.turbo,
@@ -2160,14 +2501,14 @@ def _run(
autodiff_backend=self.autodiff_backend,
cluster_manager=cluster_manager,
logger_spec=self.logger_spec,
+ use_tracing=self.use_tracing,
)
if cluster_manager is not None:
active_project = jl.seval("Base.active_project()")
if isinstance(active_project, str) and len(active_project) > 0:
- # `ClusterManagers.addprocs_slurm` launches new Julia workers via `srun`.
- # The project (environment) is propagated via `JULIA_PROJECT` rather
- # than a `--project=...` flag, so ensure it is set.
+ # Some distributed worker launchers propagate the project (environment)
+ # via `JULIA_PROJECT` rather than a `--project=...` flag.
os.environ.setdefault(
"JULIA_PROJECT", str(Path(active_project).resolve().parent)
)
@@ -2178,19 +2519,27 @@ def _run(
else:
autodiff_backend = None
- mutation_weights = SymbolicRegression.MutationWeights(
- mutate_constant=self.weight_mutate_constant,
- mutate_operator=self.weight_mutate_operator,
- mutate_feature=self.weight_mutate_feature,
- swap_operands=self.weight_swap_operands,
- rotate_tree=self.weight_rotate_tree,
- add_node=self.weight_add_node,
- insert_node=self.weight_insert_node,
- delete_node=self.weight_delete_node,
- simplify=self.weight_simplify,
- randomize=self.weight_randomize,
- do_nothing=self.weight_do_nothing,
- optimize=self.weight_optimize,
+ legacy_mutation_weights = {
+ parameter.removeprefix("weight_"): getattr(self, parameter)
+ for parameter in _LEGACY_MUTATION_PARAMETERS
+ if getattr(self, parameter) is not None
+ }
+ mutation_weights = (
+ jl_named_tuple(legacy_mutation_weights) if legacy_mutation_weights else None
+ )
+ mutations = (
+ None if self.mutations is None else convert_mutations(self.mutations)
+ )
+ default_mutations = (
+ None
+ if self.default_mutations is None
+ else convert_mutations(self.default_mutations)
+ )
+ plugins = jl_array([plugin.julia_plugin() for plugin in (self.plugins or [])])
+ default_plugins = (
+ None
+ if self.default_plugins is None
+ else jl_array([plugin.julia_plugin() for plugin in self.default_plugins])
)
# Convert operators dict to Julia format and create OperatorEnum
@@ -2201,8 +2550,12 @@ def _run(
for arity in range(1, max_arity + 1):
if arity in operators:
jl_op_list = []
- for op in operators[arity]:
- jl_op = jl.seval(op)
+ for op_index, op in enumerate(operators[arity]):
+ if type_spec_runtime is None:
+ jl_op = jl.seval(op)
+ else:
+ assert type_spec_operator_functions is not None
+ jl_op = type_spec_operator_functions[arity][op_index]
if not jl_is_function(jl_op):
raise ValueError(
f"When building operators for arity {arity}, `'{op}'` did not return a Julia function"
@@ -2213,8 +2566,17 @@ def _run(
# Empty tuple for missing arities
jl_operators_dict[arity] = ()
+ if type_spec_runtime is not None:
+ loss_type = validate_type_spec_options(
+ type_spec_runtime,
+ jl_operators_dict,
+ custom_loss,
+ )
+
complexity_mapping = (
- jl.seval(self.complexity_mapping) if self.complexity_mapping else None
+ (jl.seval(self.complexity_mapping) if self.complexity_mapping else None)
+ if type_spec_runtime is None
+ else type_spec_runtime.complexity_mapping
)
if hasattr(self, "logger_") and self.logger_ is not None and self.warm_start:
@@ -2224,23 +2586,19 @@ def _run(
self.logger_ = logger
- # Use Julia function to create OperatorEnum from Dict{Int,Tuple}
- create_operator_enum = jl.seval(
- "__sr_make_op_enum(ops_dict) = OperatorEnum([k => v for (k, v) in ops_dict]...)"
+ jl_operator_enum = SymbolicRegression.OperatorEnum(
+ tuple(jl_operators_dict.values())
)
- jl_operator_enum = create_operator_enum(jl_operators_dict)
# Build constraints dict with same structure
jl_constraints_dict = None
if any(c for c in constraints_by_arity.values() if c is not None):
- constraints_pairs = []
+ jl_constraints_dict = jl.Dict[jl.Int, jl.Vector]()
for arity in range(1, max_arity + 1):
if constraints_by_arity[arity] is not None:
- constraints_pairs.append(
- jl.Pair(arity, jl_array(constraints_by_arity[arity]))
- )
- if constraints_pairs:
- jl_constraints_dict = jl.Dict(constraints_pairs)
+ jl_constraints_dict[arity] = jl_array(constraints_by_arity[arity])
+
+ expression_spec = self._julia_expression_spec(type_spec_runtime)
# Call to Julia backend.
# See https://github.com/astroautomata/SymbolicRegression.jl/blob/master/src/OptionsStruct.jl
@@ -2251,7 +2609,7 @@ def _run(
complexity_of_constants=self.complexity_of_constants,
complexity_of_variables=complexity_of_variables,
complexity_mapping=complexity_mapping,
- expression_spec=self.expression_spec_.julia_expression_spec(),
+ expression_spec=expression_spec,
nested_constraints=nested_constraints,
elementwise_loss=custom_loss,
loss_function=custom_full_objective,
@@ -2260,16 +2618,17 @@ def _run(
maxsize=int(self.maxsize),
output_directory=_escape_filename(self.output_directory_),
npopulations=int(self.populations),
- # Determine actual batching based on "auto" mode
- batching=(self.batching if self.batching != "auto" else len(X) > 1000),
- batch_size=int(
- _get_batch_size(len(X), runtime_params.batch_size)
- if (
- self.batching == True or (self.batching == "auto" and len(X) > 1000)
- )
- else len(X)
+ batching=(
+ jl.Symbol(self.batching)
+ if isinstance(self.batching, str)
+ else self.batching
),
+ batch_size=runtime_params.batch_size,
mutation_weights=mutation_weights,
+ mutations=mutations,
+ default_mutations=default_mutations,
+ plugins=plugins,
+ default_plugins=default_plugins,
tournament_selection_p=self.tournament_selection_p,
tournament_selection_n=self.tournament_selection_n,
# These have the same name:
@@ -2310,39 +2669,43 @@ def _run(
skip_mutation_failures=self.skip_mutation_failures,
max_evals=self.max_evals,
input_stream=input_stream,
+ use_tracing=self.use_tracing,
+ tracing_file=str(self.tracing_file),
early_stop_condition=early_stop_condition,
seed=seed,
deterministic=self.deterministic,
define_helper_functions=False,
)
- self.julia_options_stream_ = jl_serialize(options)
+ serialized_options = jl_serialize(options)
+ saved_state = (
+ jl_deserialize(self.julia_state_stream_)
+ if self.warm_start and self.julia_state_stream_ is not None
+ else None
+ )
+ if self.warm_start and self.julia_options_stream_ is not None:
+ SymbolicRegression.CoreModule.check_warm_start_compatibility(
+ jl_deserialize(self.julia_options_stream_), options
+ )
# Convert data to desired precision
# This converts the data into a Julia array:
- jl_X = jl_array(np.array(X, dtype=np_dtype).T)
- if len(y.shape) == 1:
- jl_y = jl_array(np.array(y, dtype=np_dtype))
+ if type_spec_runtime is not None:
+ jl_X = type_spec_to_julia_array(type_spec_runtime, X, transpose=True)
+ jl_y = type_spec_to_julia_array(type_spec_runtime, y)
else:
- jl_y = jl_array(np.array(y, dtype=np_dtype).T)
+ jl_X = jl_numpy_array(np.array(X, dtype=np_dtype).T)
+ numeric_y = np.array(y, dtype=np_dtype)
+ jl_y = jl_numpy_array(numeric_y.T if numeric_y.ndim > 1 else numeric_y)
if weights is not None:
if len(weights.shape) == 1:
- jl_weights = jl_array(np.array(weights, dtype=np_dtype))
+ jl_weights = jl_numpy_array(np.array(weights, dtype=np_dtype))
else:
- jl_weights = jl_array(np.array(weights, dtype=np_dtype).T)
+ jl_weights = jl_numpy_array(np.array(weights, dtype=np_dtype).T)
else:
jl_weights = None
- if category is not None:
- offset_for_julia_indexing = 1
- jl_category = jl_array(
- (category + offset_for_julia_indexing).astype(np.int64)
- )
- jl_extra = jl.seval("NamedTuple{(:class,)}")((jl_category,))
- else:
- jl_extra = jl.NamedTuple()
-
if len(y.shape) > 1:
# We set these manually so that they respect Python's 0 indexing
# (by default Julia will use y1, y2...)
@@ -2352,7 +2715,13 @@ def _run(
else:
jl_y_variable_names = None
- jl_guesses = _prepare_guesses_for_julia(self.guesses, self.nout_)
+ jl_guesses = _prepare_guesses_for_julia(
+ self.guesses,
+ self.nout_,
+ expression_spec=self.expression_spec_,
+ np_dtype=np_dtype,
+ type_spec_runtime=type_spec_runtime,
+ )
# Convert worker_imports to Julia symbols
jl_worker_imports = (
@@ -2360,55 +2729,73 @@ def _run(
if self.worker_imports is not None
else None
)
-
- out = SymbolicRegression.equation_search(
- jl_X,
- jl_y,
- weights=jl_weights,
- extra=jl_extra,
- niterations=int(self.niterations),
- variable_names=jl_array([str(v) for v in self.feature_names_in_]),
- display_variable_names=jl_array(
- [str(v) for v in self.display_feature_names_in_]
- ),
- y_variable_names=jl_y_variable_names,
- X_units=jl_array(self.X_units_),
- y_units=(
- jl_array(self.y_units_)
- if isinstance(self.y_units_, list)
- else self.y_units_
- ),
- options=options,
- guesses=jl_guesses,
- numprocs=numprocs,
- parallelism=parallelism,
- saved_state=self.julia_state_,
- return_state=True,
- run_id=self.run_id_,
- addprocs_function=cluster_manager,
- heap_size_hint_in_bytes=self.heap_size_hint_in_bytes,
- worker_timeout=self.worker_timeout,
- worker_imports=jl_worker_imports,
- progress=runtime_params.progress
- and self.verbosity > 0
- and len(y.shape) == 1,
- verbosity=int(self.verbosity),
- logger=logger,
- )
+ if type_spec_runtime is not None and parallelism == "multiprocessing":
+ definition = type_spec_runtime.definition
+ assert isinstance(definition, _TypeSpecRuntimeDefinition)
+ addprocs_function = create_type_spec_addprocs_function(
+ definition,
+ cluster_manager,
+ jl_worker_imports,
+ )
+ else:
+ addprocs_function = cluster_manager
+ with _external_stop_signal_context(self) as external_stop:
+ out = SymbolicRegression.equation_search._jl_call_nogil(
+ jl_X,
+ jl_y,
+ weights=jl_weights,
+ niterations=int(self.niterations),
+ variable_names=jl_array([str(v) for v in self.feature_names_in_]),
+ display_variable_names=jl_array(
+ [str(v) for v in self.display_feature_names_in_]
+ ),
+ y_variable_names=jl_y_variable_names,
+ X_units=jl_array(self.X_units_),
+ y_units=(
+ jl_array(self.y_units_)
+ if isinstance(self.y_units_, list)
+ else self.y_units_
+ ),
+ options=options,
+ guesses=jl_guesses,
+ numprocs=numprocs,
+ parallelism=parallelism,
+ saved_state=saved_state,
+ return_state=True,
+ run_id=self.run_id_,
+ addprocs_function=addprocs_function,
+ heap_size_hint_in_bytes=self.heap_size_hint_in_bytes,
+ worker_timeout=self.worker_timeout,
+ worker_imports=jl_worker_imports,
+ progress=runtime_params.progress
+ and self.verbosity > 0
+ and len(y.shape) == 1,
+ verbosity=int(self.verbosity),
+ logger=logger,
+ external_stop=external_stop,
+ **({"loss_type": loss_type} if loss_type is not None else {}),
+ )
if self.logger_spec is not None:
self.logger_spec.write_hparams(logger, self.get_params())
if not self.warm_start:
self.logger_spec.close(logger)
- self.julia_state_stream_ = jl_serialize(out)
+ serialized_state = jl_serialize(out)
+ equations = self.get_hof(out, type_spec_runtime=type_spec_runtime)
+ if type_spec_runtime is not None:
+ assert isinstance(type_spec_runtime.definition, _TypeSpecRuntimeDefinition)
+ self._type_spec_runtime_definition_ = type_spec_runtime.definition
+ self.julia_options_stream_ = serialized_options
+ self.julia_state_stream_ = serialized_state
# Set attributes
- self.equations_ = self.get_hof(out)
+ self.equations_ = equations
ALREADY_RAN = True
return self
+ @_rollback_failed_warm_start
def fit(
self,
X,
@@ -2420,7 +2807,6 @@ def fit(
complexity_of_variables: int | float | list[int | float] | None = None,
X_units: ArrayLike[str] | None = None,
y_units: str | ArrayLike[str] | None = None,
- category: ndarray | None = None,
) -> "PySRRegressor":
"""
Search for equations to fit the dataset and store them in `self.equations_`.
@@ -2451,17 +2837,12 @@ def fit(
X_units : list[str]
A list of units for each variable in `X`. Each unit should be
a string representing a Julia expression. See DynamicQuantities.jl
- https://symbolicml.org/DynamicQuantities.jl/dev/units/ for more
+ https://ai.damtp.cam.ac.uk/dynamicquantities/dev/units for more
information.
y_units : str | list[str]
Similar to `X_units`, but as a unit for the target variable, `y`.
If `y` is a matrix, a list of units should be passed. If `X_units`
is given but `y_units` is not, then `y_units` will be arbitrary.
- category : list[int]
- If `expression_spec` is a `ParametricExpressionSpec`, then this
- argument should be a list of integers representing the category
- of each sample.
-
Returns
-------
self : object
@@ -2486,19 +2867,19 @@ def fit(
self.complexity_of_variables_ = None
self.X_units_ = None
self.y_units_ = None
+ if hasattr(self, "_type_spec_runtime_definition_"):
+ del self._type_spec_runtime_definition_
self._setup_equation_file()
self._clear_equation_file_contents()
runtime_params = self._validate_and_modify_params()
+ parallelism, numprocs = _map_parallelism_params(
+ self.parallelism,
+ self.procs,
+ getattr(self, "multithreading", None),
+ )
- if category is not None:
- assert Xresampled is None
-
- if isinstance(self.expression_spec, ParametricExpressionSpec):
- assert category is not None
-
- # TODO: Put `category` here
(
X,
y,
@@ -2519,22 +2900,9 @@ def fit(
y_units,
)
- if X.shape[0] > 50000:
- warnings.warn(
- "You are using a dataset with more than 50,000 datapoints. "
- "Symbolic regression rarely benefits from this many points - consider "
- "subsampling to 10,000 points or fewer. If you have high noise, "
- "denoise the data first rather than using more points."
- )
-
random_state = check_random_state(self.random_state) # For np random
seed = cast(int, random_state.randint(0, 2**31 - 1)) # For julia random
- if isinstance(self.expression_spec, ParametricExpressionSpec):
- parametric_expression_deprecation_warning(
- self.expression_spec.max_parameters, variable_names
- )
-
# Pre transformations (feature selection and denoising)
X, y, variable_names, complexity_of_variables, X_units, y_units = (
self._pre_transform_training_data(
@@ -2562,15 +2930,32 @@ def fit(
y,
X_units,
y_units,
+ self.type_spec is None and self._supports_export("sympy"),
)
- # Initially, just save model parameters, so that
- # it can be loaded from an early exit:
- if not self.temp_equation_file:
+ type_spec_runtime = (
+ self._load_type_spec_runtime(
+ for_fit=True,
+ operators=runtime_params.operators,
+ )
+ if self.type_spec is not None
+ else None
+ )
+ if not self.temp_equation_file and not (
+ self.warm_start and self.julia_state_stream_ is not None
+ ):
self._checkpoint()
- # Perform the search:
- self._run(X, y, runtime_params, weights=weights, seed=seed, category=category)
+ self._run(
+ X,
+ y,
+ runtime_params,
+ weights=weights,
+ seed=seed,
+ type_spec_runtime=type_spec_runtime,
+ parallelism=parallelism,
+ numprocs=numprocs,
+ )
# Then, after fit, we save again, so the pickle file contains
# the equations:
@@ -2603,8 +2988,6 @@ def predict(
self,
X,
index: int | list[int] | None = None,
- *,
- category: ndarray | None = None,
) -> ndarray:
"""
Predict y from input X using the equation chosen by `model_selection`.
@@ -2621,11 +3004,6 @@ def predict(
particular row of `self.equations_`, you may specify the index here.
For multiple output equations, you must pass a list of indices
in the same order.
- category : ndarray | None
- If `expression_spec` is a `ParametricExpressionSpec`, then this
- argument should be a list of integers representing the category
- of each sample in `X`.
-
Returns
-------
y_predicted : ndarray of shape (n_samples, nout_)
@@ -2640,70 +3018,45 @@ def predict(
check_is_fitted(
self, attributes=["selection_mask_", "feature_names_in_", "nout_"]
)
+ has_type_spec = self._has_fitted_type_spec()
best_equation = self.get_best(index=index)
- # When X is an numpy array or a pandas dataframe with a RangeIndex,
- # the self.feature_names_in_ generated during fit, for the same X,
- # will cause a warning to be thrown during _validate_data.
- # To avoid this, convert X to a dataframe, apply the selection mask,
- # and then set the column/feature_names of X to be equal to those
- # generated during fit.
- if not isinstance(X, pd.DataFrame):
- X = check_array(X)
- X = pd.DataFrame(X)
- if isinstance(X.columns, pd.RangeIndex):
- if self.selection_mask_ is not None:
- # RangeIndex enforces column order allowing columns to
- # be correctly filtered with self.selection_mask_
- X = X[X.columns[self.selection_mask_]]
- X.columns = self.feature_names_in_
-
- # During fit, we replace spaces in DataFrame column names with
- # underscores. Apply the same normalization here.
- cols_str = X.columns.astype(str)
- if cols_str.str.contains(" ").any():
- X = X.copy()
- X.columns = cols_str.str.replace(" ", "_")
- warnings.warn(
- "Spaces in DataFrame column names are not supported. "
- "Spaces have been replaced with underscores. \n"
- "Please rename the columns to valid names."
- )
-
- # Without feature information, CallableEquation/lambda_format equations
- # require that the column order of X matches that of the X used during
- # the fitting process. _validate_data removes this feature information
- # when it converts the dataframe to an np array. Thus, to ensure feature
- # order is preserved after conversion, the dataframe columns must be
- # reordered/reindexed to match those of the transformed (denoised and
- # feature selected) X in fit.
- X = X.reindex(columns=self.feature_names_in_)
- X = self._validate_data_X(X)
- if self.expression_spec_.evaluates_in_julia:
- # Julia wants the right dtype
- X = X.astype(self._get_precision_mapped_dtype(X))
-
- if category is not None:
- offset_for_julia_indexing = 1
- args: tuple = (
- jl_array((category + offset_for_julia_indexing).astype(np.int64)),
- )
+ if has_type_spec:
+ X = prepare_type_spec_prediction_data(self, X)
else:
- args = ()
+ if not isinstance(X, pd.DataFrame):
+ X = pd.DataFrame(check_array(X))
+ if isinstance(X.columns, pd.RangeIndex):
+ if self.selection_mask_ is not None:
+ X = X[X.columns[self.selection_mask_]]
+ X.columns = self.feature_names_in_
+
+ columns = X.columns.astype(str)
+ if columns.str.contains(" ").any():
+ X = X.copy()
+ X.columns = columns.str.replace(" ", "_")
+ warnings.warn(
+ "Spaces in DataFrame column names are not supported. "
+ "Spaces have been replaced with underscores. \n"
+ "Please rename the columns to valid names."
+ )
+ X = X.reindex(columns=self.feature_names_in_)
+ X = self._validate_data_X(X)
+ if self.expression_spec_.evaluates_in_julia:
+ X = X.astype(self._get_precision_mapped_dtype(X))
try:
if isinstance(best_equation, list):
assert self.nout_ > 1
return np.stack(
- [
- cast(ndarray, eq["lambda_format"](X, *args))
- for eq in best_equation
- ],
+ [cast(ndarray, eq["lambda_format"](X)) for eq in best_equation],
axis=1,
)
else:
- return cast(ndarray, best_equation["lambda_format"](X, *args))
+ return cast(ndarray, best_equation["lambda_format"](X))
except Exception as error:
+ if has_type_spec:
+ raise
raise ValueError(
"Failed to evaluate the expression. "
"If you are using a custom operator, make sure to define it in `extra_sympy_mappings`, "
@@ -2712,6 +3065,14 @@ def predict(
"You can then run `model.refresh()` to re-load the expressions."
) from error
+ def score(self, X, y, sample_weight=None):
+ if self._has_fitted_type_spec():
+ raise NotImplementedError(
+ "The R^2 `score` is not defined for models using a `type_spec`. "
+ "Evaluate predictions with a metric suited to the value type."
+ )
+ return super().score(X, y, sample_weight=sample_weight)
+
def sympy(self, index: int | list[int] | None = None):
"""
Return sympy representation of the equation(s) chosen by `model_selection`.
@@ -2730,7 +3091,7 @@ def sympy(self, index: int | list[int] | None = None):
best_equation : str, list[str] of length nout_
SymPy representation of the best equation.
"""
- if not self.expression_spec_.supports_sympy:
+ if not self._supports_export("sympy"):
raise ValueError(
f"`expression_spec={self.expression_spec_}` does not support sympy export."
)
@@ -2766,7 +3127,7 @@ def latex(
best_equation : str or list[str] of length nout_
LaTeX expression of the best equation.
"""
- if not self.expression_spec_.supports_latex:
+ if not self._supports_export("latex"):
raise ValueError(
f"`expression_spec={self.expression_spec_}` does not support latex export."
)
@@ -2803,7 +3164,7 @@ def jax(self, index=None):
Dictionary of callable jax function in "callable" key,
and jax array of parameters as "parameters" key.
"""
- if not self.expression_spec_.supports_jax:
+ if not self._supports_export("jax"):
raise ValueError(
f"`expression_spec={self.expression_spec_}` does not support jax export."
)
@@ -2839,7 +3200,7 @@ def pytorch(self, index=None):
best_equation : torch.nn.Module
PyTorch module representing the expression.
"""
- if not self.expression_spec_.supports_torch:
+ if not self._supports_export("torch"):
raise ValueError(
f"`expression_spec={self.expression_spec_}` does not support torch export."
)
@@ -2903,7 +3264,12 @@ def _postprocess_dataframe(self, df: pd.DataFrame) -> pd.DataFrame:
return df
- def get_hof(self, search_output=None) -> pd.DataFrame | list[pd.DataFrame]:
+ def get_hof(
+ self,
+ search_output=None,
+ *,
+ type_spec_runtime: _TypeSpecRuntime | None = None,
+ ) -> pd.DataFrame | list[pd.DataFrame]:
"""Get the equations from a hall of fame file or search output.
If no arguments entered, the ones used
@@ -2927,6 +3293,25 @@ def get_hof(self, search_output=None) -> pd.DataFrame | list[pd.DataFrame]:
self.equation_file_contents_ = self._read_equation_file()
_validate_export_mappings(self.extra_jax_mappings, self.extra_torch_mappings)
+ if type_spec_runtime is None and self._has_fitted_type_spec():
+ type_spec_runtime = self._load_type_spec_runtime()
+ if type_spec_runtime is not None and search_output is None:
+ search_output = jl_deserialize(self.julia_state_stream_)
+
+ def create_exports(output: pd.DataFrame, output_index: int):
+ if type_spec_runtime is not None:
+ return create_type_spec_exports(
+ type_spec_runtime,
+ output,
+ search_output,
+ output_index if self.nout_ > 1 else None,
+ )
+ return self.expression_spec_.create_exports(
+ self,
+ output,
+ search_output,
+ output_index if self.nout_ > 1 else None,
+ )
equation_file_contents = cast(List[pd.DataFrame], self.equation_file_contents_)
@@ -2941,9 +3326,7 @@ def get_hof(self, search_output=None) -> pd.DataFrame | list[pd.DataFrame]:
if self.loss_scale == "log"
else []
),
- self.expression_spec_.create_exports(
- self, output, search_output, i if self.nout_ > 1 else None
- ),
+ create_exports(output, i),
],
axis=1,
),
@@ -2983,7 +3366,7 @@ def latex_table(
latex_table_str : str
A string that will render a table in LaTeX of the equations.
"""
- if not self.expression_spec_.supports_latex:
+ if not self._supports_export("latex"):
raise ValueError(
f"`expression_spec={self.expression_spec_}` does not support latex export."
)
@@ -3070,21 +3453,15 @@ def calculate_scores(df: pd.DataFrame) -> pd.DataFrame:
)
-def _prepare_guesses_for_julia(guesses, nout) -> VectorValue | None:
- """Convert Python guesses to Julia format.
-
- Parameters
- ----------
- guesses : list[str] | list[list[str]] | list[dict[str, str]] | list[list[dict[str, str]]] | None
- Initial guesses for equations
- nout : int
- Number of output dimensions
-
- Returns
- -------
- jl_guesses: VectorValue | None
- Julia-compatible guesses array or None if no guesses provided
- """
+def _prepare_guesses_for_julia(
+ guesses,
+ nout,
+ *,
+ expression_spec: AbstractExpressionSpec,
+ np_dtype: type | None,
+ type_spec_runtime: _TypeSpecRuntime | None,
+) -> VectorValue | None:
+ """Convert Python guesses to Julia format."""
if guesses is None:
return None
@@ -3112,35 +3489,52 @@ def _prepare_guesses_for_julia(guesses, nout) -> VectorValue | None:
f"Number of guess lists ({len(g)}) must match number of outputs ({nout})"
)
+ template = (
+ expression_spec if isinstance(expression_spec, TemplateExpressionSpec) else None
+ )
+ parameter_names = set(template.parameters or {}) if template is not None else set()
+
julia_guesses = []
for output_guesses in g:
julia_output_guesses = []
for item in output_guesses:
- if isinstance(item, dict):
- # Convert dict to NamedTuple for template expressions
- julia_output_guesses.append(jl_named_tuple(item))
- else:
- # Keep strings as-is
+ if not isinstance(item, dict):
julia_output_guesses.append(item)
+ continue
+ if not parameter_names:
+ julia_output_guesses.append(jl_named_tuple(item))
+ continue
+
+ converted_item = {}
+ for name, value in item.items():
+ if name not in parameter_names:
+ converted_item[name] = value
+ continue
+ if not isinstance(value, (list, np.ndarray)):
+ raise ValueError(f"Template parameter '{name}' must be a 1D vector")
+ if type_spec_runtime is not None:
+ parameter_array = object_array_1d(value)
+ else:
+ try:
+ parameter_array = np.array(value, dtype=np_dtype, copy=True)
+ except (TypeError, ValueError) as error:
+ raise ValueError(
+ f"Template parameter '{name}' must be a 1D vector "
+ "compatible with the data precision"
+ ) from error
+ if parameter_array.ndim != 1:
+ raise ValueError(f"Template parameter '{name}' must be a 1D vector")
+ converted_item[name] = (
+ type_spec_to_julia_array(type_spec_runtime, parameter_array)
+ if type_spec_runtime is not None
+ else jl_array(parameter_array)
+ )
+ julia_output_guesses.append(jl_named_tuple(converted_item))
julia_guesses.append(jl_array(julia_output_guesses))
return jl_array(julia_guesses)
-def _get_batch_size(dataset_size: int, batch_size_param: int | None) -> int:
- """Calculate the actual batch size to use."""
- if batch_size_param is not None:
- return min(dataset_size, batch_size_param)
- elif dataset_size <= 1000:
- return dataset_size
- elif dataset_size < 5000:
- return 128
- elif dataset_size < 50000:
- return 256
- else:
- return 512
-
-
def _mutate_parameter(param_name: str, param_value):
if param_name == "batch_size" and param_value is not None and param_value < 1:
warnings.warn(
@@ -3154,15 +3548,29 @@ def _mutate_parameter(param_name: str, param_value):
and param_value == True
and "buffer" not in sys.stdout.__dir__()
):
- warnings.warn(
- "Note: it looks like you are running in Jupyter. "
- "The progress bar will be turned off."
- )
+ # The progress bar needs a real stdout buffer (e.g., not Jupyter's
+ # stream proxy), so fall back to plain printing silently.
return False
return param_value
+def _resolve_input_stream(input_stream: str | None) -> str:
+ """Map `None` to ``"stdin"`` on an interactive terminal, else ``"devnull"``.
+
+ Watching stdin for the ``'q'`` quit command (and advertising it in the
+ progress output) only makes sense when a user can actually type into
+ stdin; in notebooks, pipes, and CI it cannot work.
+ """
+ if input_stream is not None:
+ return input_stream
+ try:
+ interactive = sys.stdin.isatty()
+ except (AttributeError, ValueError):
+ interactive = False
+ return "stdin" if interactive else "devnull"
+
+
def _map_parallelism_params(
parallelism: Literal["serial", "multithreading", "multiprocessing"] | None,
procs: int | None,
diff --git a/pysr/test/__init__.py b/pysr/test/__init__.py
index 0997991fd..44cb81a75 100644
--- a/pysr/test/__init__.py
+++ b/pysr/test/__init__.py
@@ -1,6 +1,7 @@
from .test_autodiff import runtests as runtests_autodiff
from .test_cli import get_runtests as get_runtests_cli
from .test_dev import runtests as runtests_dev
+from .test_interrupt import runtests as runtests_interrupt
from .test_jax import runtests as runtests_jax
from .test_main import runtests
from .test_slurm import runtests as runtests_slurm
@@ -15,5 +16,6 @@
"get_runtests_cli",
"runtests_startup",
"runtests_dev",
+ "runtests_interrupt",
"runtests_slurm",
]
diff --git a/pysr/test/slurm_docker_cluster/config/cgroup.conf b/pysr/test/slurm_docker_cluster/config/cgroup.conf
index d6b278061..dbf068535 100644
--- a/pysr/test/slurm_docker_cluster/config/cgroup.conf
+++ b/pysr/test/slurm_docker_cluster/config/cgroup.conf
@@ -5,6 +5,8 @@
# See man slurm.conf and man cgroup.conf for further
# information on cgroup configuration parameters
#
+# CI containers use linuxproc/affinity and do not run systemd.
+CgroupPlugin=disabled
ConstrainCores=no
ConstrainDevices=no
ConstrainRAMSpace=no
diff --git a/pysr/test/test_cli.py b/pysr/test/test_cli.py
index f374ddbda..cdd82b785 100644
--- a/pysr/test/test_cli.py
+++ b/pysr/test/test_cli.py
@@ -1,5 +1,7 @@
+import os
import unittest
from textwrap import dedent
+from unittest.mock import patch
from click import testing as click_testing
@@ -52,8 +54,8 @@ def test_help_on_test(self):
Run parts of the PySR test suite.
- Choose from main, jax, torch, autodiff, cli, dev, startup, and slurm. You can
- give multiple tests, separated by commas.
+ Choose from main, jax, torch, autodiff, cli, dev, startup, slurm, and
+ interrupt. You can give multiple tests, separated by commas.
Options:
-k TEXT Filter expressions to select specific tests.
@@ -63,6 +65,51 @@ def test_help_on_test(self):
self.assertEqual(result.output.strip(), expected.strip())
self.assertEqual(result.exit_code, 0)
+ def test_test_shards_partition_filtered_duplicate_groups(self):
+ executed = []
+
+ class ShardTests(unittest.TestCase):
+ def test_ignored(self):
+ executed.append(self.id())
+
+ test_selected_a = test_ignored
+ test_selected_b = test_ignored
+ test_selected_c = test_ignored
+ test_selected_d = test_ignored
+ test_selected_e = test_ignored
+
+ loader = unittest.TestLoader()
+ selected = [
+ test.id()
+ for _ in range(2)
+ for test in loader.loadTestsFromTestCase(ShardTests)
+ if "selected" in test.id()
+ ]
+ shard_runs = []
+ with patch(
+ "pysr._cli.main.get_runtests_cli",
+ return_value=lambda **_: [ShardTests],
+ ):
+ for shard_index in range(3):
+ executed.clear()
+ with patch.dict(
+ os.environ,
+ {
+ "PYSR_TEST_SHARD_COUNT": "3",
+ "PYSR_TEST_SHARD_INDEX": str(shard_index),
+ },
+ ):
+ result = self.cli_runner.invoke(
+ pysr, ["test", "cli,cli", "-k", "selected"]
+ )
+ self.assertEqual(result.exit_code, 0, result.output)
+ shard_runs.append(executed.copy())
+
+ self.assertEqual(
+ shard_runs,
+ [selected[shard_index::3] for shard_index in range(3)],
+ )
+
def runtests(just_tests=False):
"""Run all tests in cliTest.py."""
tests = [TestCli]
diff --git a/pysr/test/test_dev_pysr.dockerfile b/pysr/test/test_dev_pysr.dockerfile
index 1a8c7e776..e757a89a2 100644
--- a/pysr/test/test_dev_pysr.dockerfile
+++ b/pysr/test/test_dev_pysr.dockerfile
@@ -4,7 +4,7 @@
ARG JLVERSION=1.11
ARG PYVERSION=3.12
-ARG BASE_IMAGE=bullseye
+ARG BASE_IMAGE=trixie
FROM julia:${JLVERSION}-${BASE_IMAGE} AS jl
FROM python:${PYVERSION}-${BASE_IMAGE}
diff --git a/pysr/test/test_interrupt.py b/pysr/test/test_interrupt.py
new file mode 100644
index 000000000..83144783b
--- /dev/null
+++ b/pysr/test/test_interrupt.py
@@ -0,0 +1,398 @@
+"""End-to-end tests for graceful SIGINT handling during a search."""
+
+import ctypes
+import os
+import signal
+import subprocess
+import sys
+import tempfile
+import textwrap
+import threading
+import time
+import unittest
+import warnings
+from pathlib import Path
+from types import SimpleNamespace
+from unittest import mock
+
+import pysr.interrupt as interrupt_module
+from pysr.sr import _resolve_input_stream as sr_resolve
+
+# Native sigaction checks are only available on POSIX.
+POSIX = os.name == "posix"
+
+
+class TestResolveInputStream(unittest.TestCase):
+ def test_default_uses_stdin_on_interactive_terminal(self):
+ with mock.patch.object(sys, "stdin") as stdin:
+ stdin.isatty.return_value = True
+ self.assertEqual(sr_resolve(None), "stdin")
+
+ def test_default_disables_stdin_watching_without_terminal(self):
+ with mock.patch.object(sys, "stdin") as stdin:
+ stdin.isatty.return_value = False
+ self.assertEqual(sr_resolve(None), "devnull")
+
+ def test_default_handles_missing_or_closed_stdin(self):
+ with mock.patch.object(sys, "stdin", None):
+ self.assertEqual(sr_resolve(None), "devnull")
+ with mock.patch.object(sys, "stdin") as stdin:
+ stdin.isatty.side_effect = ValueError("I/O operation on closed file")
+ self.assertEqual(sr_resolve(None), "devnull")
+
+ def test_explicit_values_pass_through(self):
+ for value in ("stdin", "devnull", "Main.my_stream"):
+ self.assertEqual(sr_resolve(value), value)
+
+
+@unittest.skipUnless(POSIX, "SIGINT signal contexts are POSIX-only")
+class TestExternalStopSignalContext(unittest.TestCase):
+ def setUp(self):
+ self.native_handler_before = self._native_handler()
+ self.python_before = signal.getsignal(signal.SIGINT)
+ self.stop_fds = []
+ backend = SimpleNamespace(ExternalStop=self._external_stop)
+ self.patches = (
+ mock.patch.dict(os.environ, {"PYTHON_JULIACALL_HANDLE_SIGNALS": "yes"}),
+ mock.patch.object(interrupt_module, "SymbolicRegression", backend),
+ )
+ for patcher in self.patches:
+ patcher.start()
+ self.addCleanup(patcher.stop)
+
+ def tearDown(self):
+ self.assertEqual(signal.getsignal(signal.SIGINT), self.python_before)
+ self.assertEqual(self._native_handler(), self.native_handler_before)
+
+ def _external_stop(self, read_fd, trigger):
+ self.stop_fds.append(read_fd)
+ return SimpleNamespace(fd=read_fd, trigger=trigger)
+
+ @staticmethod
+ def _native_handler():
+ # Compare only `sa_handler`. It is the whole of what this test
+ # asserts, and as the first member on both Linux and Darwin it needs
+ # no platform-specific struct offsets to read.
+ class Action(ctypes.Union):
+ _fields_ = [
+ ("handler", ctypes.c_void_p),
+ ("alignment", ctypes.c_longdouble),
+ ("storage", ctypes.c_ubyte * 1024),
+ ]
+
+ libc = ctypes.CDLL(None, use_errno=True)
+ libc.sigaction.argtypes = [ctypes.c_int, ctypes.c_void_p, ctypes.c_void_p]
+ libc.sigaction.restype = ctypes.c_int
+ action = Action()
+ result = libc.sigaction(signal.SIGINT, None, ctypes.byref(action))
+ if result != 0:
+ errno = ctypes.get_errno()
+ raise OSError(errno, os.strerror(errno))
+ return action.handler
+
+ @staticmethod
+ def _current_wakeup_fd():
+ current = signal.set_wakeup_fd(-1)
+ signal.set_wakeup_fd(current)
+ return current
+
+ def _assert_stop_fds_closed(self):
+ for fd in self.stop_fds:
+ with self.assertRaises(OSError):
+ os.fstat(fd)
+
+ def test_cleanup_after_normal_return_restores_existing_wakeup_fd(self):
+ wakeup_read, wakeup_write = os.pipe()
+ os.set_blocking(wakeup_write, False)
+ previous_wakeup = signal.set_wakeup_fd(wakeup_write)
+ try:
+ model = SimpleNamespace()
+ with interrupt_module._external_stop_signal_context(model) as external_stop:
+ self.assertEqual(external_stop.fd, self.stop_fds[0])
+ self.assertEqual(external_stop.trigger, signal.SIGINT)
+ self.assertFalse(os.get_blocking(external_stop.fd))
+ self.assertFalse(os.get_blocking(self._current_wakeup_fd()))
+ self.assertEqual(self._current_wakeup_fd(), wakeup_write)
+ self.assertFalse(model.interrupted_)
+ self._assert_stop_fds_closed()
+ finally:
+ signal.set_wakeup_fd(previous_wakeup)
+ os.close(wakeup_write)
+ os.close(wakeup_read)
+
+ def test_cleanup_after_search_error(self):
+ model = SimpleNamespace()
+ with self.assertRaisesRegex(RuntimeError, "search failed"):
+ with interrupt_module._external_stop_signal_context(model):
+ raise RuntimeError("search failed")
+ self._assert_stop_fds_closed()
+
+ def test_cleanup_continues_after_cleanup_failure(self):
+ real_set_wakeup_fd = signal.set_wakeup_fd
+ calls = []
+
+ def fail_after_restoring(fd):
+ result = real_set_wakeup_fd(fd)
+ calls.append(fd)
+ if len(calls) == 2:
+ raise RuntimeError("injected cleanup failure")
+ return result
+
+ model = SimpleNamespace()
+ with mock.patch.object(signal, "set_wakeup_fd", fail_after_restoring):
+ with self.assertRaisesRegex(RuntimeError, "injected cleanup failure"):
+ with interrupt_module._external_stop_signal_context(model):
+ pass
+ self._assert_stop_fds_closed()
+
+ def test_signals_disabled_does_not_arm(self):
+ model = SimpleNamespace(interrupted_=True)
+ with (
+ mock.patch.dict(os.environ, {"PYTHON_JULIACALL_HANDLE_SIGNALS": "no"}),
+ mock.patch.object(os, "pipe") as pipe,
+ ):
+ with interrupt_module._external_stop_signal_context(model) as external_stop:
+ self.assertIsNone(external_stop)
+ pipe.assert_not_called()
+ self.assertFalse(model.interrupted_)
+
+ def test_interrupted_return_sets_attribute_and_warns(self):
+ model = SimpleNamespace()
+ with warnings.catch_warnings(record=True) as caught:
+ warnings.simplefilter("always")
+ with interrupt_module._external_stop_signal_context(model):
+ handler = signal.getsignal(signal.SIGINT)
+ handler(signal.SIGINT, None)
+ self.assertTrue(model.interrupted_)
+ self.assertTrue(any("partial" in str(item.message).lower() for item in caught))
+
+
+# The child asserts partial results, restored native SIGINT disposition on
+# POSIX, and that a second fit in the same process works.
+CHILD_SCRIPT = textwrap.dedent("""
+ import ctypes
+ import os
+ import signal
+ import threading
+ import time
+ import warnings
+
+ import numpy as np
+
+ from pysr import PySRRegressor
+
+ if os.name == "posix":
+ libc = ctypes.CDLL(None)
+
+ def native_handler():
+ # Compare only `sa_handler`: the first member on both platforms, and
+ # all this check needs.
+ storage = (ctypes.c_char * 512)()
+ libc.sigaction(int(signal.SIGINT), None, storage)
+ return ctypes.cast(storage, ctypes.POINTER(ctypes.c_void_p)).contents.value
+
+ before = native_handler()
+
+ def send_interrupt():
+ if os.name == "posix":
+ os.kill(os.getpid(), signal.SIGINT)
+ elif not ctypes.windll.kernel32.GenerateConsoleCtrlEvent(0, 0):
+ raise ctypes.WinError()
+
+ rstate = np.random.RandomState(0)
+ X = rstate.randn(150, 2)
+ y = X[:, 0] * X[:, 1]
+ warmup = PySRRegressor(
+ niterations=1,
+ populations=3,
+ verbosity=0,
+ progress=False,
+ temp_equation_file=True,
+ )
+ warmup.fit(X[:20], y[:20])
+
+ model = PySRRegressor(
+ niterations=1_000_000, # only an interrupt can end this fit
+ populations=8,
+ verbosity=0,
+ progress=False,
+ temp_equation_file=True,
+ )
+ unarmed_handler = signal.getsignal(signal.SIGINT)
+
+ def interrupt_once_armed():
+ # The fit installs its own SIGINT handler once the search is armed; the
+ # extra wait lets the search run long enough to have partial results.
+ while signal.getsignal(signal.SIGINT) is unarmed_handler:
+ time.sleep(0.1)
+ time.sleep(5.0)
+ send_interrupt()
+
+ threading.Thread(target=interrupt_once_armed, daemon=True).start()
+ with warnings.catch_warnings(record=True) as caught:
+ warnings.simplefilter("always")
+ model.fit(X, y)
+ assert model.equations_ is not None
+ assert len(model.equations_) > 0, "search returned no partial results"
+ assert model.interrupted_ is True
+ assert any("partial" in str(item.message).lower() for item in caught)
+ print(f"INTERRUPTED_OK:{len(model.equations_)}", flush=True)
+
+ if os.name == "posix":
+ print(f"HANDLER_RESTORED:{native_handler() == before}", flush=True)
+
+ model2 = PySRRegressor(
+ niterations=2,
+ populations=8,
+ verbosity=0,
+ progress=False,
+ temp_equation_file=True,
+ )
+ model2.fit(X, y)
+ assert model2.equations_ is not None
+ assert model2.interrupted_ is False
+ print("SECOND_FIT_OK", flush=True)
+ """)
+
+# Wait this long after Jupyter's SEARCHING marker so the fit has armed the
+# cooperative handler before the single user interrupt.
+FIRST_SIGNAL_DELAY = 15.0
+TOTAL_TIMEOUT = 600.0
+
+
+class TestSubprocessInterrupt(unittest.TestCase):
+ def test_sigint_returns_partial_results_and_restores_state(self):
+ with tempfile.TemporaryDirectory() as tmpdir:
+ script = Path(tmpdir) / "child.py"
+ script.write_text(CHILD_SCRIPT)
+ p = subprocess.Popen(
+ [sys.executable, str(script)],
+ cwd=tmpdir,
+ stdout=subprocess.PIPE,
+ stderr=subprocess.STDOUT,
+ text=True,
+ creationflags=0 if POSIX else subprocess.CREATE_NEW_CONSOLE,
+ )
+ watchdog = threading.Timer(TOTAL_TIMEOUT, p.kill)
+ watchdog.start()
+ lines = []
+ try:
+ for line in p.stdout:
+ lines.append(line.strip())
+ p.wait()
+ finally:
+ watchdog.cancel()
+ if p.poll() is None:
+ p.kill()
+
+ output = "\n".join(lines)
+ self.assertEqual(p.returncode, 0, f"child failed:\n{output}")
+ self.assertIn("INTERRUPTED_OK:", output)
+ if POSIX:
+ self.assertIn("HANDLER_RESTORED:True", output)
+ self.assertIn("SECOND_FIT_OK", output)
+
+
+JUPYTER_FIT_CELL = textwrap.dedent("""
+ import numpy as np
+ from pysr import PySRRegressor
+
+ rstate = np.random.RandomState(0)
+ X = rstate.randn(150, 2)
+ y = X[:, 0] * X[:, 1]
+
+ # Compile the search pipeline before the timed window opens, so that the
+ # interrupt lands on a running search rather than on Julia's warm-up.
+ warmup = PySRRegressor(
+ niterations=1,
+ populations=3,
+ verbosity=0,
+ progress=False,
+ temp_equation_file=True,
+ )
+ warmup.fit(X[:20], y[:20])
+
+ model = PySRRegressor(
+ niterations=1_000_000,
+ populations=8,
+ verbosity=0,
+ progress=False,
+ temp_equation_file=True,
+ )
+ print("SEARCHING", flush=True)
+ model.fit(X, y)
+ assert model.interrupted_ is True
+ assert len(model.equations_) > 0, "search returned no partial results"
+ print("INTERRUPTED_OK", len(model.equations_), flush=True)
+ """)
+
+
+class TestJupyterInterrupt(unittest.TestCase):
+ """Drive a real ipykernel and interrupt it exactly like the Jupyter UI."""
+
+ def test_kernel_interrupt_stops_search_and_kernel_survives(self):
+ try:
+ from jupyter_client.manager import start_new_kernel
+ except ImportError:
+ self.skipTest("jupyter_client is not installed")
+
+ km, kc = start_new_kernel()
+ try:
+ msg_id = kc.execute(JUPYTER_FIT_CELL)
+ self._await_stream(kc, "SEARCHING", timeout=TOTAL_TIMEOUT)
+
+ # Interrupt once, exactly like one click on Jupyter's interrupt action.
+ time.sleep(FIRST_SIGNAL_DELAY)
+ km.interrupt_kernel()
+ reply = self._await_reply(kc, msg_id, timeout=TOTAL_TIMEOUT)
+ assert reply is not None, "kernel never returned from fit"
+ # The fit returns normally with partial results: no error status.
+ self.assertEqual(reply["content"]["status"], "ok")
+
+ # The kernel must remain fully usable afterwards.
+ reply = kc.execute_interactive(
+ "print('ALIVE', len(model.equations_), flush=True)",
+ timeout=60,
+ )
+ self.assertEqual(reply["content"]["status"], "ok")
+ finally:
+ kc.stop_channels()
+ km.shutdown_kernel(now=True)
+
+ # -- helpers -----------------------------------------------------------
+
+ def _await_stream(self, kc, text, timeout):
+ deadline = time.monotonic() + timeout
+ while time.monotonic() < deadline:
+ msg = kc.get_iopub_msg(timeout=deadline - time.monotonic())
+ if msg["msg_type"] == "stream" and text in msg["content"]["text"]:
+ return
+ raise AssertionError(f"never saw {text!r} on iopub")
+
+ def _await_reply(self, kc, msg_id, timeout):
+ deadline = time.monotonic() + timeout
+ while time.monotonic() < deadline:
+ try:
+ msg = kc.get_shell_msg(timeout=max(0.1, deadline - time.monotonic()))
+ except Exception:
+ return None
+ if msg["parent_header"].get("msg_id") == msg_id:
+ return msg
+ return None
+
+
+def runtests(just_tests=False):
+ tests = [
+ TestResolveInputStream,
+ TestExternalStopSignalContext,
+ TestSubprocessInterrupt,
+ TestJupyterInterrupt,
+ ]
+ if just_tests:
+ return tests
+ suite = unittest.TestSuite()
+ loader = unittest.TestLoader()
+ for test in tests:
+ suite.addTests(loader.loadTestsFromTestCase(test))
+ runner = unittest.TextTestRunner()
+ return runner.run(suite)
diff --git a/pysr/test/test_jax.py b/pysr/test/test_jax.py
index d261afcf5..45db9f3c7 100644
--- a/pysr/test/test_jax.py
+++ b/pysr/test/test_jax.py
@@ -1,5 +1,7 @@
+import subprocess
+import sys
+import tempfile
import unittest
-from functools import partial
from pathlib import Path
import numpy as np
@@ -49,7 +51,7 @@ def test_pipeline_pandas(self):
for fname in ["hall_of_fame.csv.bak", "hall_of_fame.csv"]:
equations["Complexity Loss Equation".split(" ")].to_csv(
- Path(model.output_directory_) / model.run_id_ / fname
+ Path(model.output_directory_) / model.run_id_ / fname, index=False
)
model.refresh(run_directory=str(Path(model.output_directory_) / model.run_id_))
@@ -77,7 +79,7 @@ def test_pipeline(self):
for fname in ["hall_of_fame.csv.bak", "hall_of_fame.csv"]:
equations["Complexity Loss Equation".split(" ")].to_csv(
- Path(model.output_directory_) / model.run_id_ / fname
+ Path(model.output_directory_) / model.run_id_ / fname, index=False
)
model.refresh(run_directory=str(Path(model.output_directory_) / model.run_id_))
@@ -125,29 +127,53 @@ def cos_approx(x):
return 1 - (x**2) / 2 + (x**4) / 24 + (x**6) / 720
sp_cos_approx = sympy.Function("cos_approx")
-
y = X["k15"] ** 2 + 2 * cos_approx(X["k20"])
- model = PySRRegressor(
- progress=False,
- unary_operators=["cos_approx(x) = 1 - x^2 / 2 + x^4 / 24 + x^6 / 720"],
- select_k_features=3,
- maxsize=10,
- early_stop_condition=1e-5,
- extra_sympy_mappings={"cos_approx": sp_cos_approx},
- extra_jax_mappings={
- sp_cos_approx: "(lambda x: 1 - x**2 / 2 + x**4 / 24 + x**6 / 720)"
- },
- random_state=0,
- deterministic=True,
- parallelism="serial",
- )
- np.random.seed(0)
- model.fit(X.values, y.values)
- f, parameters = model.jax().values()
- jax_prediction = partial(f, parameters=parameters)
- jax_output = jax_prediction(X.values)
- np.testing.assert_almost_equal(y.values, jax_output, decimal=3)
+ with tempfile.TemporaryDirectory() as directory:
+ model = PySRRegressor(
+ progress=False,
+ unary_operators=["cos_approx(x) = 1 - x^2 / 2 + x^4 / 24 + x^6 / 720"],
+ select_k_features=3,
+ maxsize=10,
+ early_stop_condition=1e-5,
+ extra_sympy_mappings={"cos_approx": sp_cos_approx},
+ extra_jax_mappings={
+ sp_cos_approx: "(lambda x: 1 - x**2 / 2 + x**4 / 24 + x**6 / 720)"
+ },
+ random_state=0,
+ deterministic=True,
+ parallelism="serial",
+ output_directory=directory,
+ run_id="custom-jax-checkpoint",
+ temp_equation_file=False,
+ )
+ np.random.seed(0)
+ model.fit(X.values, y.values)
+ self.assertTrue(model.show_pickle_warnings_)
+
+ input_path = Path(directory) / "input.npy"
+ output_path = Path(directory) / "output.npy"
+ np.save(input_path, X.values)
+ run_directory = Path(directory) / "custom-jax-checkpoint"
+ code = f"""
+import numpy as np
+from pysr import PySRRegressor
+model = PySRRegressor.from_file(run_directory={str(run_directory)!r})
+f, parameters = model.jax().values()
+np.save({str(output_path)!r}, np.asarray(f(np.load({str(input_path)!r}), parameters)))
+"""
+ result = subprocess.run(
+ [sys.executable, "-c", code],
+ check=False,
+ capture_output=True,
+ text=True,
+ )
+ self.assertEqual(result.returncode, 0, result.stderr)
+ symbols = [sympy.Symbol(name) for name in model.feature_names_in_]
+ reference = sympy.lambdify(
+ symbols, model.sympy(), modules=[{"cos_approx": cos_approx}, "numpy"]
+ )(*X.values[:, model.selection_mask_].T)
+ np.testing.assert_almost_equal(reference, np.load(output_path), decimal=3)
def runtests(just_tests=False):
diff --git a/pysr/test/test_main.py b/pysr/test/test_main.py
index b703ff752..453eaa50a 100644
--- a/pysr/test/test_main.py
+++ b/pysr/test/test_main.py
@@ -1,13 +1,19 @@
import functools
import importlib
+import json
import os
import pickle as pkl
import platform
-import re
+import shutil
+import subprocess
+import sys
import tempfile
+import threading
+import time
import traceback
import unittest
import warnings
+from dataclasses import fields, is_dataclass
from pathlib import Path
from textwrap import dedent
from unittest import mock
@@ -24,7 +30,6 @@
estimator_checks_generator = functools.partial(check_estimator, generate_only=True)
from pysr import (
- ParametricExpressionSpec,
PySRRegressor,
TemplateExpressionSpec,
TensorBoardLoggerSpec,
@@ -34,11 +39,11 @@
)
from pysr.export_latex import sympy2latex
from pysr.export_sympy import pysr2sympy
-from pysr.expression_specs import parametric_expression_deprecation_warning
from pysr.feature_selection import _handle_feature_selection, run_feature_selection
-from pysr.julia_helpers import init_julia
+from pysr.julia_helpers import _load_cluster_manager, init_julia
from pysr.sr import (
_check_assertions,
+ _create_julia_operators_and_loss_functions,
_process_constraints,
_suggest_keywords,
_validate_elementwise_loss,
@@ -53,6 +58,7 @@
DEFAULT_POPULATIONS,
skip_if_beartype,
)
+from .test_type_specs import TestTypeSpecs
# Disables local saving:
os.environ["SYMBOLIC_REGRESSION_IS_TESTING"] = os.environ.get(
@@ -353,6 +359,17 @@ def test_elementwise_loss_float32_fit_accepts_strictly_typed_loss(self):
def test_validation_helpers_skip_nonfunction(self):
_validate_elementwise_loss(jl.seval("1.0"), has_weights=False)
+ def test_elementwise_loss_accepts_lossfunctions_object(self):
+ _, custom_loss, _, _ = _create_julia_operators_and_loss_functions(
+ operators={2: ["+", "-", "/", "*"]},
+ extra_sympy_mappings=None,
+ supports_sympy=True,
+ elementwise_loss="LPDistLoss{3}()",
+ loss_function=None,
+ loss_function_expression=None,
+ )
+ self.assertTrue(jl.applicable(custom_loss, 1.0, 1.0))
+
def test_validate_export_mappings_typechecks(self):
with self.assertRaises(ValueError):
_validate_export_mappings({"a": 1}, None)
@@ -621,6 +638,10 @@ def test_warm_start_set_at_init(self):
regressor = PySRRegressor(warm_start=True, max_evals=10)
regressor.fit(self.X, y)
+ regressor.set_params(binary_operators=["+", "*", "-", "/"])
+ with self.assertRaisesRegex(JuliaError, "Warm start incompatible.*operators"):
+ regressor.fit(self.X, y)
+
def test_noisy_builtin_variable_names(self):
y = self.X[:, [0, 1]] ** 2 + self.rstate.randn(self.X.shape[0], 1) * 0.05
model = PySRRegressor(
@@ -832,7 +853,9 @@ def test_template_expressions_and_custom_complexity(self):
# Create model with template that includes the missing sin operator
model = PySRRegressor(
expression_spec=TemplateExpressionSpec(
- ["f"], "sin_of_f((; f), (x, y)) = sin(f(x, y))"
+ combine="sin(f(x, y))",
+ expressions=["f"],
+ variable_names=["x", "y"],
),
binary_operators=["+", "-", "*", "/"],
unary_operators=[], # No sin operator!
@@ -879,6 +902,112 @@ def test_template_expressions_and_custom_complexity(self):
with self.assertRaises(ValueError):
model.latex_table()
+ def test_template_expressions_reload_in_fresh_process(self):
+ X = self.rstate.uniform(-1, 1, (30, 2))
+ y = np.sin(X[:, 0] + X[:, 1])
+ with tempfile.TemporaryDirectory() as directory:
+ model = PySRRegressor(
+ expression_spec=TemplateExpressionSpec(
+ combine="sin(f(x, y))",
+ expressions=["f"],
+ variable_names=["x", "y"],
+ ),
+ binary_operators=["+", "-", "*"],
+ unary_operators=[],
+ maxsize=10,
+ **{
+ **self.default_test_kwargs,
+ "temp_equation_file": False,
+ },
+ output_directory=directory,
+ run_id="template-checkpoint",
+ early_stop_condition=1e-10,
+ )
+ model.fit(X, y)
+ expected = model.predict(X)
+ run_directory = Path(directory) / "template-checkpoint"
+ code = f"""
+import json
+import numpy as np
+from pysr import PySRRegressor
+model = PySRRegressor.from_file(run_directory={str(run_directory)!r})
+X = np.array({X.tolist()!r})
+print(json.dumps({{
+ "columns": list(model.equations_.columns),
+ "prediction": model.predict(X).tolist(),
+}}))
+"""
+ result = subprocess.run(
+ [sys.executable, "-c", code],
+ check=False,
+ capture_output=True,
+ text=True,
+ )
+
+ self.assertEqual(result.returncode, 0, result.stderr)
+ payload = json.loads(result.stdout.strip().splitlines()[-1])
+ self.assertIn("lambda_format", payload["columns"])
+ np.testing.assert_allclose(payload["prediction"], expected, atol=1e-10)
+
+ loaded = pkl.loads(pkl.dumps(model))
+ self.assertIsNotNone(loaded.julia_state_)
+
+ def test_template_expression_pickle_roundtrip_in_fresh_process(self):
+ X = self.rstate.uniform(-1, 1, (30, 2))
+ y = np.sin(X[:, 0] + X[:, 1])
+ with tempfile.TemporaryDirectory() as directory:
+ model = PySRRegressor(
+ expression_spec=TemplateExpressionSpec(
+ combine="sin(f(x, y))",
+ expressions=["f"],
+ variable_names=["x", "y"],
+ ),
+ binary_operators=["+", "-", "*"],
+ unary_operators=[],
+ maxsize=10,
+ **{
+ **self.default_test_kwargs,
+ "temp_equation_file": False,
+ },
+ output_directory=directory,
+ run_id="template-pickle",
+ early_stop_condition=1e-10,
+ )
+ model.fit(X, y)
+ expected = model.predict(X)
+
+ pickle_path = Path(directory) / "template-model.pkl"
+ with open(pickle_path, "wb") as f:
+ pkl.dump(model, f)
+ # Delete the run directory: the pickle alone must be enough to
+ # rebuild the Julia-backed equation columns.
+ shutil.rmtree(Path(directory) / "template-pickle")
+
+ code = f"""
+import json
+import pickle
+import numpy as np
+
+with open({str(pickle_path)!r}, "rb") as f:
+ model = pickle.load(f)
+X = np.array({X.tolist()!r})
+prediction = model.predict(X)
+print(json.dumps({{
+ "columns": list(model.equations_.columns),
+ "prediction": prediction.tolist(),
+}}))
+"""
+ result = subprocess.run(
+ [sys.executable, "-c", code],
+ check=False,
+ capture_output=True,
+ text=True,
+ )
+ self.assertEqual(result.returncode, 0, result.stderr)
+ payload = json.loads(result.stdout.strip().splitlines()[-1])
+ self.assertIn("lambda_format", payload["columns"])
+ np.testing.assert_allclose(payload["prediction"], expected, atol=1e-10)
+
def test_template_expression_with_parameters_multiout(self):
# Create random data
X_continuous = self.rstate.uniform(-1, 1, (100, 2))
@@ -897,7 +1026,7 @@ def test_template_expression_with_parameters_multiout(self):
model = PySRRegressor(
**self.default_test_kwargs,
expression_spec=TemplateExpressionSpec(
- "p[class] * x1^2 + f(x2)",
+ combine="p[class] * x1^2 + f(x2)",
expressions=["f"],
parameters={"p": 3},
variable_names=["x1", "x2", "class"],
@@ -927,61 +1056,6 @@ def test_template_expression_with_parameters_multiout(self):
test_mse = np.mean((y_test - y_pred) ** 2)
self.assertLess(test_mse, 1e-5)
- def test_parametric_expression(self):
- # Create data with two classes
- n_points = 100
- X = self.rstate.uniform(-3, 3, (n_points, 2)) # x1, x2
- category = self.rstate.randint(0, 3, n_points) # class (0 or 1)
-
- # True parameters for each class
- P1 = [0.1, 1.5, -5.2] # phase shift for each class
- P2 = [3.2, 0.5, 1.2] # offset for each class
-
- # Ground truth: 2*cos(x2 + P1[class]) + x1^2 - P2[class]
- y = np.array(
- [
- 2 * np.cos(x2 + P1[c]) + x1**2 - P2[c]
- for x1, x2, c in zip(X[:, 0], X[:, 1], category)
- ]
- )
-
- model = PySRRegressor(
- expression_spec=ParametricExpressionSpec(max_parameters=2),
- binary_operators=["+", "*", "/", "-"],
- unary_operators=["cos", "exp"],
- maxsize=20,
- early_stop_condition="stop_if(loss, complexity) = loss < 1e-4 && complexity <= 14",
- **self.default_test_kwargs,
- )
-
- model.fit(X, y, category=category)
-
- # Test on new data points
- X_test = self.rstate.uniform(-6, 6, (10, 2))
- category_test = self.rstate.randint(0, 3, 10)
-
- y_test = np.array(
- [
- 2 * np.cos(x2 + P1[c]) + x1**2 - P2[c]
- for x1, x2, c in zip(X_test[:, 0], X_test[:, 1], category_test)
- ]
- )
-
- y_test_pred = model.predict(X_test, category=category_test)
- test_mse = np.mean((y_test - y_test_pred) ** 2)
- self.assertLess(test_mse, 1e-3)
-
- with self.assertRaises(ValueError):
- model.sympy()
- with self.assertRaises(ValueError):
- model.latex()
- with self.assertRaises(ValueError):
- model.jax()
- with self.assertRaises(ValueError):
- model.pytorch()
- with self.assertRaises(ValueError):
- model.latex_table()
-
def test_tensorboard_logger(self):
if platform.system() == "Windows":
@@ -1043,6 +1117,37 @@ def test_tensorboard_logger(self):
# Verify model still works as expected
self.assertLessEqual(model.get_best()["loss"], 1e-4)
+ def test_tracing_writes_jsonl(self):
+ """Tracing should write the search and each iteration as JSON lines."""
+ X = self.rstate.randn(10, 2)
+ y = X[:, 0]
+ kwargs = dict(
+ niterations=2,
+ populations=2,
+ population_size=20,
+ ncycles_per_iteration=DEFAULT_NCYCLES,
+ progress=False,
+ temp_equation_file=True,
+ parallelism="serial",
+ deterministic=True,
+ random_state=0,
+ )
+ with tempfile.TemporaryDirectory() as tmpdir:
+ trace_file = Path(tmpdir) / "trace.jsonl"
+
+ PySRRegressor(**kwargs, use_tracing=True, tracing_file=trace_file).fit(X, y)
+
+ records = [json.loads(line) for line in trace_file.read_text().splitlines()]
+ self.assertEqual(records[0]["record_type"], "search")
+ iterations = [r for r in records if r["record_type"] == "iteration"]
+ self.assertGreater(len(iterations), 0)
+ self.assertIn("tree", iterations[0]["members"][0])
+ self.assertIn("mutations", iterations[0])
+
+ trace_file.unlink()
+ PySRRegressor(**kwargs).fit(X, y)
+ self.assertFalse(trace_file.exists())
+
def test_negative_losses(self):
X = self.rstate.rand(100, 3) * 20.0
variance = X[:, 0] * 0.1
@@ -1191,6 +1296,29 @@ def test_template_expression_guesses(self):
# Check that a good fit was found
self.assertTrue(any(model.equations_["loss"] < 1e-10))
+ def test_template_expression_guess_parameters(self):
+ X = self.rstate.randn(100, 1)
+ y = 5.0 * X[:, 0] ** 2 + 10.0 * X[:, 0] + 0.8
+ parameter_values = np.array([5.0, 10.0, 0.8], dtype=np.float64)
+ original_values = parameter_values.copy()
+ model = PySRRegressor(
+ expression_spec=TemplateExpressionSpec(
+ expressions=["f"],
+ combine="p[1] * f(x) + p[2] * x + p[3]",
+ variable_names=["x"],
+ parameters={"p": 3, "unused": 1},
+ ),
+ guesses=[{"f": "#1 * #1", "p": parameter_values}],
+ should_optimize_constants=False,
+ precision=32,
+ **self.default_test_kwargs,
+ )
+ model.fit(X, y, variable_names=["x"])
+
+ np.testing.assert_allclose(model.predict(X), y, rtol=1e-6, atol=1e-6)
+ np.testing.assert_array_equal(parameter_values, original_values)
+ self.assertEqual(parameter_values.dtype, np.float64)
+
def test_multi_output_template_expression_guesses(self):
X = self.rstate.randn(100, 2)
Y = np.column_stack([X[:, 0] + X[:, 1], X[:, 0] - X[:, 1]])
@@ -1210,6 +1338,77 @@ def test_multi_output_template_expression_guesses(self):
for i, eqs_df in enumerate(model.equations_):
self.assertTrue(any(eqs_df["loss"] < 1e-10))
+ def test_template_parameter_guesses_replace_on_warm_start(self):
+ X = self.rstate.randn(100, 1)
+ y = 3.0 * X[:, 0] + 4.0
+ template = TemplateExpressionSpec(
+ expressions=["f"],
+ combine="p[1] * f(x) + p[2]",
+ variable_names=["x"],
+ parameters={"p": 2},
+ )
+ model = PySRRegressor(
+ expression_spec=template,
+ guesses=[{"f": "#1", "p": [1.0, 0.0]}],
+ should_optimize_constants=False,
+ **self.default_test_kwargs,
+ )
+ model.fit(X, y, variable_names=["x"])
+ model.set_params(guesses=[{"f": "#1", "p": [3.0, 4.0]}], warm_start=True)
+ model.fit(X, y, variable_names=["x"])
+
+ np.testing.assert_allclose(model.predict(X), y, rtol=1e-6, atol=1e-6)
+
+ def test_multi_output_template_parameter_guesses(self):
+ X = self.rstate.randn(100, 1)
+ Y = np.column_stack([2.0 * X[:, 0] + 1.0, -3.0 * X[:, 0] + 0.5])
+ model = PySRRegressor(
+ expression_spec=TemplateExpressionSpec(
+ expressions=["f"],
+ combine="p[1] * f(x) + p[2]",
+ variable_names=["x"],
+ parameters={"p": 2},
+ ),
+ guesses=[
+ [{"f": "#1", "p": [2.0, 1.0]}],
+ [{"f": "#1", "p": [-3.0, 0.5]}],
+ ],
+ should_optimize_constants=False,
+ **self.default_test_kwargs,
+ )
+ model.fit(X, Y, variable_names=["x"])
+
+ prediction = model.predict(X)
+ self.assertEqual(prediction.shape, Y.shape)
+ np.testing.assert_allclose(prediction, Y, rtol=1e-6, atol=1e-6)
+
+ def test_invalid_template_parameter_guesses(self):
+ X = self.rstate.randn(10, 1)
+ y = X[:, 0]
+ template = TemplateExpressionSpec(
+ expressions=["f"],
+ combine="p[1] * f(x) + p[2]",
+ variable_names=["x"],
+ parameters={"p": 2},
+ )
+ cases = [
+ ({"f": "#1", "missing": "#1", "p": [1.0, 2.0]}, JuliaError),
+ ({"p": [1.0, 2.0]}, JuliaError),
+ ({"f": ["#1"], "p": [1.0, 2.0]}, JuliaError),
+ ({"f": "#1", "p": [1.0]}, JuliaError),
+ ({"f": "#1", "p": 1.0}, ValueError),
+ ({"f": "#1", "p": [[1.0, 2.0]]}, ValueError),
+ ]
+ for guess, error in cases:
+ with self.subTest(guess=guess):
+ model = PySRRegressor(
+ expression_spec=template,
+ guesses=[guess],
+ **self.default_test_kwargs,
+ )
+ with self.assertRaises(error):
+ model.fit(X, y, variable_names=["x"])
+
def test_invalid_multi_output_format_guesses(self):
X = self.rstate.randn(100, 2)
Y = np.column_stack([X[:, 0], X[:, 1]])
@@ -1343,14 +1542,16 @@ def manually_create_model(equations, feature_names=None):
Path(output_directory)
/ run_id
/ f"hall_of_fame_output{i+1}.csv.bak"
- )
+ ),
+ index=False,
)
else:
model.nout_ = 1
model.selection_mask_ = None
model.feature_names_in_ = np.array(feature_names, dtype=object)
equations["complexity loss equation".split(" ")].to_csv(
- str(Path(output_directory) / run_id / "hall_of_fame.csv.bak")
+ str(Path(output_directory) / run_id / "hall_of_fame.csv.bak"),
+ index=False,
)
model.refresh()
@@ -1437,6 +1638,142 @@ def test_feature_selection_handler(self):
class TestMiscellaneous(unittest.TestCase):
"""Test miscellaneous functions."""
+ def test_search_releases_the_gil(self):
+ """Other Python threads must keep running while Julia searches."""
+ gaps = []
+ stop = threading.Event()
+
+ def sample_gaps():
+ previous = time.monotonic()
+ while not stop.is_set():
+ time.sleep(0.01)
+ now = time.monotonic()
+ gaps.append(now - previous)
+ previous = now
+
+ rstate = np.random.RandomState(0)
+ X = rstate.randn(100, 2)
+ y = X[:, 0] * X[:, 1]
+ model = PySRRegressor(
+ niterations=40,
+ populations=8,
+ verbosity=0,
+ progress=False,
+ temp_equation_file=True,
+ )
+
+ watcher = threading.Thread(target=sample_gaps, daemon=True)
+ watcher.start()
+ started = time.monotonic()
+ try:
+ model.fit(X, y)
+ finally:
+ elapsed = time.monotonic() - started
+ stop.set()
+ watcher.join(timeout=5)
+
+ # Holding the GIL across the search would leave one gap spanning it.
+ self.assertLess(max(gaps), elapsed / 4)
+
+ def test_loss_function_can_call_back_into_python(self):
+ """A Julia loss_function shim can drive a Python objective mid-search."""
+ calls = {"n": 0}
+
+ def python_objective(tree, dataset, options):
+ calls["n"] += 1
+ prediction, completed = jl.SymbolicRegression.eval_tree_array(
+ tree, dataset.X, options
+ )
+ if not completed:
+ return float("inf")
+ prediction = np.asarray(prediction)
+ target = np.asarray(dataset.y)
+ return float(np.mean((prediction - target) ** 2))
+
+ jl.pysr_test_python_objective = python_objective
+ jl.seval("""
+ using PythonCall
+
+ function pysr_test_python_objective_shim(
+ tree, dataset::Dataset{T,L}, options
+ )::L where {T,L}
+ PythonCall.GIL.@lock begin
+ return pyconvert(
+ L, pysr_test_python_objective(tree, dataset, options)
+ )
+ end
+ end
+ """)
+
+ X = np.arange(8.0).reshape(-1, 1)
+ y = 2.0 * X[:, 0]
+ model = PySRRegressor(
+ binary_operators=["+"],
+ unary_operators=[],
+ loss_function="pysr_test_python_objective_shim",
+ niterations=1,
+ populations=1,
+ population_size=20,
+ ncycles_per_iteration=1,
+ maxsize=7,
+ should_optimize_constants=False,
+ parallelism="serial",
+ deterministic=True,
+ random_state=0,
+ precision=64,
+ progress=False,
+ verbosity=0,
+ temp_equation_file=True,
+ )
+ model.fit(X, y)
+
+ self.assertGreater(calls["n"], 0)
+ self.assertTrue(np.all(np.isfinite(model.equations_["loss"])))
+
+ def test_unfitted_julia_streams_are_none(self):
+ model = PySRRegressor()
+ self.assertIsNone(model.julia_options_)
+ self.assertIsNone(model.julia_state_)
+
+ def test_inline_operator_must_be_a_function(self):
+ from pysr.sr import _maybe_create_inline_operators
+
+ with self.assertRaisesRegex(ValueError, "must evaluate to a Julia function"):
+ _maybe_create_inline_operators({1: ["identity(1)"]}, None, False)
+
+ def test_variable_names_must_match_feature_count(self):
+ # sklearn's own length check fires first during `fit`; exercise the
+ # backstop in `_check_assertions` directly.
+ from pysr.sr import _check_assertions
+
+ X = np.random.randn(10, 2)
+ with self.assertRaisesRegex(ValueError, "one name per feature"):
+ _check_assertions(X, True, ["a"], None, None, X[:, 0], None, None, False)
+
+ def test_equations_with_quoted_constants_are_read(self):
+ """Equations holding string constants arrive with doubled quotes."""
+ output_directory = tempfile.mkdtemp()
+ run_id = "quoted"
+ os.makedirs(Path(output_directory) / run_id)
+ equation = 'f = "a"; p = ["a", "b"]'
+ with open(
+ Path(output_directory) / run_id / "hall_of_fame.csv", "w", encoding="utf-8"
+ ) as f:
+ f.write("Complexity,Loss,Equation\n")
+ f.write('1,0.0,"""a"""\n')
+ f.write('3,0.0,"f = ""a""; p = [""a"", ""b""]"\n')
+
+ model = PySRRegressor()
+ model.nout_ = 1
+ model.output_directory_ = output_directory
+ model.run_id_ = run_id
+
+ equations = model._read_equation_file()[0]
+
+ self.assertEqual(list(equations["complexity"]), [1, 3])
+ self.assertEqual(list(equations["loss"]), [0.0, 0.0])
+ self.assertEqual(list(equations["equation"]), ['"a"', equation])
+
def test_pickle_inv_sympy_expression(self):
"""Test that sympy expressions with the inv operator can be pickled and unpickled correctly."""
expr_str = "inv(x0) + x1"
@@ -1518,6 +1855,104 @@ def test_pickle_with_temp_equation_file(self):
y_predictions2 = model2.predict(X)
np.testing.assert_array_almost_equal(y_predictions, y_predictions2)
+ def test_checkpoint_schema(self):
+ state = PySRRegressor().__getstate__()
+ self.assertEqual(state["_checkpoint_schema_version"], 3)
+
+ schema_two_state = state.copy()
+ schema_two_state.pop("type_spec")
+ schema_two_model = PySRRegressor.__new__(PySRRegressor)
+ schema_two_model.__setstate__(schema_two_state)
+ self.assertIsNone(schema_two_model.type_spec)
+
+ missing_schema_state = state.copy()
+ missing_schema_state.pop("_checkpoint_schema_version")
+ missing_schema_model = PySRRegressor.__new__(PySRRegressor)
+ with self.assertRaisesRegex(ValueError, "Unsupported PySR checkpoint schema"):
+ missing_schema_model.__setstate__(missing_schema_state)
+
+ for schema_version in (1, 2, 4):
+ incompatible_state = state.copy()
+ incompatible_state["_checkpoint_schema_version"] = schema_version
+
+ model = PySRRegressor.__new__(PySRRegressor)
+ with self.assertRaisesRegex(
+ ValueError, "Unsupported PySR checkpoint schema"
+ ):
+ model.__setstate__(incompatible_state)
+
+ def test_checkpoint_dump_failure_preserves_existing_file(self):
+ def fail_after_partial_write(_, checkpoint_file):
+ checkpoint_file.write(b"partial")
+ raise TypeError("cannot pickle")
+
+ with tempfile.TemporaryDirectory() as directory:
+ model = PySRRegressor()
+ model.output_directory_ = directory
+ model.run_id_ = "failed-checkpoint"
+ model.show_pickle_warnings_ = False
+ checkpoint = model.get_pkl_filename()
+ checkpoint.write_bytes(b"previous checkpoint")
+
+ with mock.patch("pysr.sr.pkl.dump", side_effect=fail_after_partial_write):
+ model._checkpoint()
+
+ self.assertEqual(checkpoint.read_bytes(), b"previous checkpoint")
+ self.assertFalse(model.show_pickle_warnings_)
+ self.assertEqual(list(checkpoint.parent.glob("*.tmp")), [])
+
+ def test_checkpoint_dump_failure_leaves_no_file(self):
+ def fail_after_partial_write(_, checkpoint_file):
+ checkpoint_file.write(b"partial")
+ raise TypeError("cannot pickle")
+
+ with tempfile.TemporaryDirectory() as directory:
+ model = PySRRegressor()
+ model.output_directory_ = directory
+ model.run_id_ = "failed-checkpoint"
+ checkpoint = model.get_pkl_filename()
+
+ with mock.patch("pysr.sr.pkl.dump", side_effect=fail_after_partial_write):
+ model._checkpoint()
+
+ self.assertFalse(checkpoint.exists())
+ self.assertTrue(model.show_pickle_warnings_)
+ self.assertEqual(list(checkpoint.parent.glob("*.tmp")), [])
+
+ def test_checkpoint_cleanup_failure_restores_warnings(self):
+ with tempfile.TemporaryDirectory() as directory:
+ model = PySRRegressor()
+ model.output_directory_ = directory
+ model.run_id_ = "failed-cleanup"
+ model.show_pickle_warnings_ = True
+ checkpoint = model.get_pkl_filename()
+
+ with mock.patch(
+ "pathlib.Path.unlink", side_effect=OSError("cannot unlink")
+ ):
+ model._checkpoint()
+
+ self.assertTrue(checkpoint.exists())
+ self.assertTrue(model.show_pickle_warnings_)
+
+ def test_checkpoint_propagates_publication_failure(self):
+ with tempfile.TemporaryDirectory() as directory:
+ model = PySRRegressor()
+ model.output_directory_ = directory
+ model.run_id_ = "unpublishable-checkpoint"
+ model.show_pickle_warnings_ = True
+ checkpoint = model.get_pkl_filename()
+
+ with mock.patch(
+ "pysr.sr.os.replace", side_effect=OSError("read-only directory")
+ ):
+ with self.assertRaises(OSError):
+ model._checkpoint()
+
+ self.assertFalse(checkpoint.exists())
+ self.assertEqual(list(checkpoint.parent.glob("*.tmp")), [])
+ self.assertTrue(model.show_pickle_warnings_)
+
def test_scikit_learn_compatibility(self):
"""Test PySRRegressor compatibility with scikit-learn."""
model = PySRRegressor(
@@ -1578,7 +2013,7 @@ def test_invalid_batch_size_corrects_and_warns(self):
self.assertTrue(any("batch_size" in str(w.message) for w in caught))
- def test_progress_disabled_when_stdout_lacks_buffer(self):
+ def test_progress_silently_disabled_when_stdout_lacks_buffer(self):
fake_stdout = type(
"FakeStdout", (), {"write": lambda self, *_args, **_kwargs: None}
)()
@@ -1596,7 +2031,286 @@ def test_progress_disabled_when_stdout_lacks_buffer(self):
y = X[:, 0]
model.fit(X, y)
- self.assertTrue(any("progress bar" in str(w.message) for w in caught))
+ self.assertFalse(any("progress bar" in str(w.message) for w in caught))
+
+ def test_builtin_mutation_and_plugin_configs(self):
+ from pysr import (
+ AdaptiveMutationWeightsPlugin,
+ BacksolveMutation,
+ ConstantMutation,
+ MutationBurstPlugin,
+ OperatorMutation,
+ SymbolicRegression,
+ )
+
+ mutation_name = jl.seval("m -> string(nameof(typeof(m)))")
+ for mutation, expected_name in (
+ (OperatorMutation(), "OperatorMutation"),
+ (BacksolveMutation(), "BacksolveMutation"),
+ ):
+ julia_mutation = mutation.julia_mutation()
+ self.assertEqual(str(mutation_name(julia_mutation)), expected_name)
+
+ backsolve_mutation = BacksolveMutation(
+ max_library_size=321,
+ max_terms=4,
+ min_improvement=0.02,
+ node_attempts=5,
+ ).julia_mutation()
+ self.assertEqual(backsolve_mutation.max_library_size, 321)
+ self.assertEqual(backsolve_mutation.max_terms, 4)
+ self.assertEqual(backsolve_mutation.min_improvement, 0.02)
+ self.assertEqual(backsolve_mutation.node_attempts, 5)
+ constant_mutation = ConstantMutation().julia_mutation()
+ backend_default = SymbolicRegression.ConstantMutation()
+ self.assertEqual(
+ constant_mutation.perturbation_factor,
+ backend_default.perturbation_factor,
+ )
+ self.assertEqual(
+ constant_mutation.probability_negate,
+ backend_default.probability_negate,
+ )
+
+ plugin_name = jl.seval("p -> string(nameof(typeof(p)))")
+ for plugin, expected_name in (
+ (
+ AdaptiveMutationWeightsPlugin(reward="loss"),
+ "AdaptiveMutationWeightsPlugin",
+ ),
+ (MutationBurstPlugin(retry_attempts=2), "MutationBurstPlugin"),
+ ):
+ self.assertEqual(str(plugin_name(plugin.julia_plugin())), expected_name)
+
+ def test_default_operators_and_plugins_match_backend(self):
+ from pysr import SymbolicRegression
+
+ model = PySRRegressor()
+ backend_options = SymbolicRegression.Options()
+
+ backend_binary_operators = jl.seval(
+ "options -> string.(options.operators.ops[2])"
+ )(backend_options)
+ self.assertListEqual(
+ model._validate_and_modify_params().operators[2],
+ [str(operator) for operator in backend_binary_operators],
+ )
+
+ plugin_name = jl.seval("p -> string(nameof(typeof(p)))")
+ plugins = {
+ str(plugin_name(plugin)): plugin for plugin in backend_options.plugins
+ }
+
+ parsimony = plugins["AdaptiveParsimonyPlugin"]
+ self.assertEqual(model.use_frequency, parsimony.mutation_acceptance)
+ self.assertEqual(model.use_frequency_in_tournament, parsimony.tournament)
+
+ def test_mutation_and_plugin_configuration(self):
+ from pysr import (
+ AbstractPlugin,
+ AdaptiveParsimonyPlugin,
+ BacksolveMutation,
+ ConstantMutation,
+ RandomizeMutation,
+ SimulatedAnnealingPlugin,
+ SymbolicRegression,
+ )
+
+ jl.seval("""
+ if !isdefined(Main, :PySRPluginMutationTest)
+ @eval module PySRPluginMutationTest
+ using SymbolicRegression
+ struct PluginMutation <: SymbolicRegression.AbstractMutation end
+ struct Plugin <: SymbolicRegression.AbstractPlugin end
+ SymbolicRegression.plugin_mutations(::Plugin) =
+ (PluginMutation() => 0.25,)
+ end
+ end
+ """)
+
+ class PluginConfig(AbstractPlugin):
+ def julia_plugin(self):
+ return jl.seval("PySRPluginMutationTest.Plugin()")
+
+ X = np.arange(12, dtype=np.float32).reshape(6, 2)
+ y = X[:, 0]
+ model = PySRRegressor(
+ binary_operators=["+"],
+ unary_operators=[],
+ niterations=0,
+ populations=1,
+ population_size=4,
+ tournament_selection_n=2,
+ progress=False,
+ temp_equation_file=True,
+ perturbation_factor=0.7,
+ probability_negate_constant=0.8,
+ mutations={
+ BacksolveMutation(max_library_size=123): 0.02,
+ ConstantMutation(perturbation_factor=0.2): 0.5,
+ RandomizeMutation(): 0.4,
+ },
+ plugins=[
+ SimulatedAnnealingPlugin(alpha=0.4),
+ AdaptiveParsimonyPlugin(tournament=False, mutation_acceptance=False),
+ PluginConfig(),
+ ],
+ )
+ model.fit(X, y)
+
+ options = model.julia_options_
+ mutation_name = jl.seval("p -> string(nameof(typeof(first(p))))")
+ mutation_by_name = {
+ str(mutation_name(pair)): pair for pair in options.mutations
+ }
+ self.assertEqual(options.crossover_probability, 0.2)
+ expected_weights = {
+ str(mutation_name(pair)): float(jl.last(pair))
+ for pair in SymbolicRegression.Options().mutations
+ }
+ expected_weights.update(
+ ConstantMutation=0.5,
+ RandomizeMutation=0.4,
+ BacksolveMutation=0.02,
+ PluginMutation=0.25,
+ )
+ self.assertSetEqual(set(mutation_by_name), set(expected_weights))
+ for name, weight in expected_weights.items():
+ self.assertEqual(float(jl.last(mutation_by_name[name])), weight)
+ constant_mutation = jl.first(mutation_by_name["ConstantMutation"])
+ self.assertEqual(constant_mutation.perturbation_factor, 0.2)
+ self.assertEqual(
+ constant_mutation.probability_negate,
+ SymbolicRegression.ConstantMutation().probability_negate,
+ )
+ self.assertEqual(
+ jl.first(mutation_by_name["BacksolveMutation"]).max_library_size,
+ 123,
+ )
+
+ plugin_name = jl.seval("p -> string(nameof(typeof(p)))")
+ plugins = {str(plugin_name(plugin)): plugin for plugin in options.plugins}
+ self.assertSetEqual(
+ set(plugins),
+ {
+ "SimulatedAnnealingPlugin",
+ "AdaptiveParsimonyPlugin",
+ "AdaptiveMutationWeightsPlugin",
+ "Plugin",
+ },
+ )
+ self.assertEqual(plugins["SimulatedAnnealingPlugin"].alpha, 0.4)
+ self.assertFalse(plugins["AdaptiveParsimonyPlugin"].tournament)
+ self.assertFalse(plugins["AdaptiveParsimonyPlugin"].mutation_acceptance)
+
+ def test_default_mutation_and_plugin_configuration(self):
+ from pysr import (
+ BacksolveMutation,
+ MutationBurstPlugin,
+ OperatorMutation,
+ SimulatedAnnealingPlugin,
+ )
+
+ X = np.arange(12, dtype=np.float32).reshape(6, 2)
+ y = X[:, 0]
+ model = PySRRegressor(
+ binary_operators=["+"],
+ unary_operators=[],
+ niterations=0,
+ populations=1,
+ population_size=4,
+ tournament_selection_n=2,
+ progress=False,
+ temp_equation_file=True,
+ default_mutations={OperatorMutation(): 0.7},
+ mutations={BacksolveMutation(max_library_size=123): 0.2},
+ default_plugins=[MutationBurstPlugin(retry_attempts=2)],
+ plugins=[SimulatedAnnealingPlugin(alpha=0.4)],
+ )
+ model.fit(X, y)
+
+ mutation_name = jl.seval("p -> string(nameof(typeof(first(p))))")
+ weights = {
+ str(mutation_name(pair)): float(jl.last(pair))
+ for pair in model.julia_options_.mutations
+ }
+ self.assertDictEqual(
+ weights,
+ {"OperatorMutation": 0.7, "BacksolveMutation": 0.2},
+ )
+
+ plugin_name = jl.seval("p -> string(nameof(typeof(p)))")
+ plugins = {
+ str(plugin_name(plugin)): plugin for plugin in model.julia_options_.plugins
+ }
+ self.assertSetEqual(
+ set(plugins), {"MutationBurstPlugin", "SimulatedAnnealingPlugin"}
+ )
+ self.assertEqual(plugins["MutationBurstPlugin"].retry_attempts, 2)
+ self.assertEqual(plugins["SimulatedAnnealingPlugin"].alpha, 0.4)
+
+ def test_legacy_mutation_weight_override(self):
+ from pysr import SymbolicRegression
+
+ X = np.arange(12, dtype=np.float32).reshape(6, 2)
+ y = X[:, 0]
+ model = PySRRegressor(
+ binary_operators=["+"],
+ unary_operators=[],
+ niterations=0,
+ populations=1,
+ population_size=4,
+ tournament_selection_n=2,
+ progress=False,
+ temp_equation_file=True,
+ perturbation_factor=0.7,
+ probability_negate_constant=0.8,
+ weight_randomize=0.3,
+ weight_backsolve=0.2,
+ )
+ model.fit(X, y)
+
+ mutation_name = jl.seval("p -> string(nameof(typeof(first(p))))")
+ weights = {
+ str(mutation_name(pair)): float(jl.last(pair))
+ for pair in model.julia_options_.mutations
+ }
+ self.assertEqual(weights["RandomizeMutation"], 0.3)
+ self.assertEqual(weights["BacksolveMutation"], 0.2)
+ backend_weights = {
+ str(mutation_name(pair)): float(jl.last(pair))
+ for pair in SymbolicRegression.default_mutations()
+ }
+ self.assertEqual(weights["AddNodeMutation"], backend_weights["AddNodeMutation"])
+ constant_mutation = next(
+ jl.first(pair)
+ for pair in model.julia_options_.mutations
+ if str(mutation_name(pair)) == "ConstantMutation"
+ )
+ self.assertEqual(constant_mutation.perturbation_factor, 0.7)
+ self.assertEqual(constant_mutation.probability_negate, 0.8)
+
+ def test_mutation_interfaces_are_mutually_exclusive(self):
+ from sklearn.base import clone
+
+ from pysr import OperatorMutation
+
+ model = PySRRegressor()
+ cloned_model = clone(model)
+ legacy_parameters = filter(
+ lambda parameter: parameter.startswith("weight_"), model.get_params()
+ )
+ for parameter in legacy_parameters:
+ self.assertIsNone(getattr(model, parameter))
+ self.assertIsNone(getattr(cloned_model, parameter))
+
+ for new_interface in (
+ {"default_mutations": {OperatorMutation(): 0.7}},
+ {"mutations": {OperatorMutation(): 0.7}},
+ ):
+ conflicting_model = PySRRegressor(weight_add_node=2.47, **new_interface)
+ with self.assertRaisesRegex(ValueError, "Cannot combine legacy"):
+ conflicting_model._validate_and_modify_params()
def test_param_groupings(self):
"""Test that param_groupings are complete"""
@@ -1621,54 +2335,18 @@ def test_param_groupings(self):
# Check the sets are equal:
self.assertSetEqual(set(params), set(regressor_params))
- def test_parametric_deprecation_warning(self):
- """Test that the helpful warning message is displayed."""
- pattern = re.compile(
- r"ParametricExpressionSpec is deprecated.*TemplateExpressionSpec.*"
- r"max_parameters=2.*"
- r"variable_names=\[\"alpha\", \"beta\"\].*"
- r"expressions=\[\"f\"\].*"
- r"variable_names=\[\"alpha\", \"beta\", \"category\"\].*"
- r"parameters=\{\s*\"p1\": n_categories,\s*\"p2\": n_categories\s*\}.*"
- r"combine=\"f\(alpha, beta, p1\[category\], p2\[category\]\)\"",
- flags=re.S,
- )
-
- with self.assertWarnsRegex(FutureWarning, pattern):
- parametric_expression_deprecation_warning(
- max_parameters=2,
- variable_names=["alpha", "beta"],
- )
-
def test_load_all_packages(self):
"""Test we can load all packages at once."""
load_all_packages()
self.assertTrue(jl.seval("ClusterManagers isa Module"))
+ self.assertTrue(jl.seval("SlurmClusterManager isa Module"))
- def test_get_batch_size(self):
- """Test the _get_batch_size function."""
- from pysr.sr import _get_batch_size
-
- # Test None (auto) mode with different dataset sizes
- self.assertEqual(_get_batch_size(500, None), 500)
- self.assertEqual(_get_batch_size(999, None), 999)
- self.assertEqual(_get_batch_size(1000, None), 1000)
- self.assertEqual(_get_batch_size(1001, None), 128)
- self.assertEqual(_get_batch_size(1500, None), 128)
- self.assertEqual(_get_batch_size(4999, None), 128)
- self.assertEqual(_get_batch_size(5000, None), 256)
- self.assertEqual(_get_batch_size(10000, None), 256)
- self.assertEqual(_get_batch_size(49999, None), 256)
- self.assertEqual(_get_batch_size(50000, None), 512)
- self.assertEqual(_get_batch_size(100000, None), 512)
-
- # Test explicit batch_size
- self.assertEqual(_get_batch_size(1000, 64), 64)
- self.assertEqual(_get_batch_size(1000, 2000), 1000) # Capped at dataset size
- self.assertEqual(_get_batch_size(50, 100), 50) # Capped at dataset size
+ with mock.patch.dict(os.environ, {"SLURM_JOB_ID": "1", "SLURM_NTASKS": "2"}):
+ with self.assertRaisesRegex(JuliaError, "allocation has 2 tasks"):
+ _load_cluster_manager("slurm")(1)
def test_batching_auto(self):
- """Test that batching='auto' works correctly."""
+ """Test that batching configuration is passed to SymbolicRegression.jl."""
model = PySRRegressor()
self.assertEqual(model.batching, "auto")
@@ -1676,11 +2354,13 @@ def test_batching_auto(self):
y_small = np.random.randn(100)
model = PySRRegressor(batching="auto", niterations=0)
model.fit(X_small, y_small)
+ self.assertEqual(model.julia_options_.batching, jl.Symbol("auto"))
+ self.assertIsNone(model.julia_options_.batch_size)
- X_large = np.random.randn(1001, 2)
- y_large = np.random.randn(1001)
- model2 = PySRRegressor(batching="auto", niterations=0)
- model2.fit(X_large, y_large)
+ explicit = PySRRegressor(batching=True, batch_size=64, niterations=0)
+ explicit.fit(X_small, y_small)
+ self.assertTrue(explicit.julia_options_.batching)
+ self.assertEqual(explicit.julia_options_.batch_size, 64)
def test_batch_size_negative_warning(self):
"""Test that batch_size < 1 gives a warning for integers only."""
@@ -1752,17 +2432,6 @@ def test_from_file_requires_operator_configuration(self):
self.assertIn("must provide either `operators`", str(cm.exception))
- def test_size_warning(self):
- """Ensure that a warning is given for a large input size."""
- model = PySRRegressor()
- X = np.random.randn(50001, 2)
- y = np.random.randn(50001)
- with warnings.catch_warnings():
- warnings.simplefilter("error")
- with self.assertRaises(Exception) as context:
- model.fit(X, y)
- self.assertIn("more than 50,000", str(context.exception))
-
def test_deterministic_warnings(self):
"""Ensure that warnings are given for determinism"""
model = PySRRegressor(random_state=0)
@@ -2150,6 +2819,7 @@ def test_unit_checks(self):
y,
X_units,
y_units,
+ supports_sympy=False,
)
invalid_units = [
(np.ones((10, 2)), np.ones(10), ["m/s", "s", "s^2"], None),
@@ -2165,6 +2835,7 @@ def test_unit_checks(self):
y,
X_units,
y_units,
+ supports_sympy=False,
)
def test_unit_propagation(self):
@@ -2246,22 +2917,20 @@ def test_process_constraints_swaps_multiplication_constraints(self):
class TestTemplateExpressionSpec(unittest.TestCase):
- def test_num_features_symbol_keys(self):
- # ponytail: one check — dict keys must reach Julia as Symbols
- spec = TemplateExpressionSpec(
- ["f", "g"],
- "combine(fs, vars) = fs.f(vars[1], vars[2]) + fs.g(vars[3])",
- {"f": 2, "g": 1},
- )
- options = spec.julia_expression_options()
- names = jl.seval("x -> propertynames(x.structure.num_features)")(options)
- self.assertEqual(names, (jl.Symbol("f"), jl.Symbol("g")))
-
def _check_macro_str(self, spec, expected_str):
self.assertEqual(
spec._template_macro_str().strip(), dedent(expected_str).strip()
)
+ def test_is_dataclass(self):
+ spec = TemplateExpressionSpec("f(x)", ["f"], ["x"])
+
+ self.assertTrue(is_dataclass(spec))
+ self.assertEqual(
+ [field.name for field in fields(spec)],
+ ["combine", "expressions", "variable_names", "parameters"],
+ )
+
def test_single_expression_no_params_single_variable(self):
spec = TemplateExpressionSpec(
combine="f(x)", expressions=["f"], variable_names=["x"]
@@ -2382,6 +3051,7 @@ def runtests(just_tests=False):
TestLaTeXTable,
TestDimensionalConstraints,
TestGuesses,
+ TestTypeSpecs,
]
if just_tests:
return test_cases
diff --git a/pysr/test/test_nb.ipynb b/pysr/test/test_nb.ipynb
index 0a2050a2a..f2b93382c 100644
--- a/pysr/test/test_nb.ipynb
+++ b/pysr/test/test_nb.ipynb
@@ -4,19 +4,12 @@
"cell_type": "code",
"execution_count": 1,
"metadata": {},
- "outputs": [
- {
- "name": "stdout",
- "output_type": "stream",
- "text": [
- "Detected Jupyter notebook. Loading juliacall extension. Set `PYSR_AUTOLOAD_EXTENSIONS=no` to disable.\n"
- ]
- }
- ],
+ "outputs": [],
"source": [
"# NBVAL_IGNORE_OUTPUT\n",
"import numpy as np\n",
- "from pysr import PySRRegressor, jl"
+ "from pysr import PySRRegressor, jl\n",
+ "%load_ext juliacall"
]
},
{
@@ -35,7 +28,7 @@
"source": [
"%%julia\n",
"\n",
- "# Automatically activates Julia magic\n",
+ "# Julia magic comes from the juliacall extension loaded above\n",
"\n",
"x = 1\n",
"println(x + 2)"
diff --git a/pysr/test/test_slurm.py b/pysr/test/test_slurm.py
index bda248c91..eae0679a8 100644
--- a/pysr/test/test_slurm.py
+++ b/pysr/test/test_slurm.py
@@ -154,9 +154,10 @@ def test_pysr_slurm_cluster_manager(self):
"#SBATCH --job-name=pysr-slurm-test",
"#SBATCH --partition=normal",
"#SBATCH --nodes=2",
- "#SBATCH --ntasks=4",
+ "#SBATCH --ntasks-per-node=2",
"#SBATCH --time=40:00",
"set -euo pipefail",
+ "export JULIA_DEBUG=SlurmClusterManager",
"python3 - <<'PY'",
"import numpy as np",
"from pysr import PySRRegressor",
@@ -164,11 +165,11 @@ def test_pysr_slurm_cluster_manager(self):
"y = X[:, 0] + 1.0",
"model = PySRRegressor(",
" niterations=3,",
- " populations=3,",
+ " populations=4,",
" progress=False,",
" temp_equation_file=True,",
" parallelism='multiprocessing',",
- " procs=2,",
+ " procs=4,",
" cluster_manager='slurm',",
" verbosity=0,",
")",
@@ -275,6 +276,15 @@ def test_pysr_slurm_cluster_manager(self):
env=self.compose_env,
)
self.assertEqual(output.returncode, 0, msg=output.stdout)
+ worker_hosts = re.findall(
+ r"Worker \d+ ready on host ([^,]+), port \d+", output.stdout
+ )
+ self.assertEqual(len(worker_hosts), 4, msg=output.stdout)
+ self.assertEqual(
+ sorted(worker_hosts.count(host) for host in set(worker_hosts)),
+ [2, 2],
+ msg=output.stdout,
+ )
self.assertIn("PYSR_SLURM_OK", output.stdout)
diff --git a/pysr/test/test_startup.py b/pysr/test/test_startup.py
index 8bd81cd93..9a0686d42 100644
--- a/pysr/test/test_startup.py
+++ b/pysr/test/test_startup.py
@@ -131,6 +131,86 @@ def test_bad_startup_options(self):
)
self.assertIn(warning_test["msg"], result.stderr.decode())
+ def test_autoload_extension_env_precedence(self):
+ autoload_key = "PYTHON_JULIACALL_AUTOLOAD_IPYTHON_EXTENSION"
+ deprecated_key = "PYSR_AUTOLOAD_EXTENSIONS"
+ cases = [
+ # (extra env, expected value after `import pysr`)
+ ({}, "no"),
+ ({autoload_key: "yes"}, "yes"),
+ ({deprecated_key: "yes"}, "yes"),
+ # An explicitly set juliacall var beats the deprecated alias:
+ ({autoload_key: "no", deprecated_key: "yes"}, "no"),
+ ]
+ for extra_env, expected in cases:
+ env = {
+ k: v
+ for k, v in os.environ.items()
+ if k not in (autoload_key, deprecated_key)
+ }
+ env.update(extra_env)
+ result = subprocess.run(
+ [
+ sys.executable,
+ "-c",
+ f'import pysr, os; print("autoload=" + os.environ["{autoload_key}"])',
+ ],
+ stdout=subprocess.PIPE,
+ stderr=subprocess.PIPE,
+ env=env,
+ )
+ self.assertEqual(result.returncode, 0, result.stderr.decode())
+ self.assertIn(f"autoload={expected}", result.stdout.decode())
+ deprecation_expected = deprecated_key in extra_env
+ self.assertEqual(
+ "PYSR_AUTOLOAD_EXTENSIONS is deprecated" in result.stderr.decode(),
+ deprecation_expected,
+ )
+
+ def test_juliacall_0935_distributed_worker_uses_current_python(self):
+ import_orders = (
+ "import pysr\nfrom pysr import jl",
+ "import juliacall\nimport pysr\nfrom pysr import jl",
+ )
+ worker_code = r"""
+ using Distributed
+ worker = only(addprocs(1, exeflags="--threads=1"))
+ try
+ fetch(Distributed.remotecall_eval(Main, worker, :(using PythonCall)))
+ @fetchfrom worker (
+ pyconvert(String, pyimport("sys").executable),
+ pyconvert(Int, pybuiltins.sum([1, 2, 3])),
+ )
+ finally
+ rmprocs(worker)
+ end
+ """
+ env = os.environ.copy()
+ env.pop("JULIA_PYTHONCALL_EXE", None)
+ env.pop("JULIA_PYTHONCALL_EXECUTABLE", None)
+ env["JULIA_CONDAPKG_OFFLINE"] = "yes"
+ env["PYTHON_JULIACALL_THREADS"] = "1"
+ env["JULIA_NUM_THREADS"] = "1"
+
+ for import_order in import_orders:
+ with self.subTest(import_order=import_order):
+ code = (
+ "import sys\n"
+ f"{import_order}\n"
+ f"worker_python, total = jl.seval({worker_code!r})\n"
+ "assert worker_python == sys.executable, "
+ "(worker_python, sys.executable)\n"
+ "assert total == 6, total\n"
+ )
+ result = subprocess.run(
+ [sys.executable, "-c", code],
+ capture_output=True,
+ text=True,
+ env=env,
+ timeout=300,
+ )
+ self.assertEqual(result.returncode, 0, result.stdout + result.stderr)
+
def test_notebook(self):
if platform.system() == "Windows":
self.skipTest("Notebook test incompatible with Windows")
diff --git a/pysr/test/test_torch.py b/pysr/test/test_torch.py
index e66e7d960..0d0933d0f 100644
--- a/pysr/test/test_torch.py
+++ b/pysr/test/test_torch.py
@@ -51,7 +51,7 @@ def test_pipeline_pandas(self):
for fname in ["hall_of_fame.csv.bak", "hall_of_fame.csv"]:
equations["Complexity Loss Equation".split(" ")].to_csv(
- Path(model.output_directory_) / model.run_id_ / fname
+ Path(model.output_directory_) / model.run_id_ / fname, index=False
)
model.refresh(run_directory=str(Path(model.output_directory_) / model.run_id_))
@@ -85,7 +85,7 @@ def test_pipeline(self):
for fname in ["hall_of_fame.csv.bak", "hall_of_fame.csv"]:
equations["Complexity Loss Equation".split(" ")].to_csv(
- Path(model.output_directory_) / model.run_id_ / fname
+ Path(model.output_directory_) / model.run_id_ / fname, index=False
)
model.refresh(run_directory=str(Path(model.output_directory_) / model.run_id_))
@@ -138,7 +138,7 @@ def test_custom_operator(self):
for fname in ["hall_of_fame.csv.bak", "hall_of_fame.csv"]:
equations["Complexity Loss Equation".split(" ")].to_csv(
- Path(model.output_directory_) / model.run_id_ / fname
+ Path(model.output_directory_) / model.run_id_ / fname, index=False
)
MyCustomOperator = sympy.Function("mycustomoperator")
@@ -194,7 +194,7 @@ def test_issue_656(self):
def test_constant_arguments(self):
# Test that functions with constant arguments work correctly
- # Regression test for https://github.com/MilesCranmer/PySR/issues/656
+ # Regression test for https://github.com/astroautomata/PySR/issues/656
test_cases = [
(pysr.export_sympy.pysr2sympy("sqrt(2)"), np.sqrt(2)),
(sympy.exp(2), np.exp(2)),
@@ -246,8 +246,7 @@ def cos_approx(x):
torch_output = torch_module(self.torch.tensor(X.values)).detach().numpy()
- np.testing.assert_almost_equal(y.values, np_output, decimal=3)
- np.testing.assert_almost_equal(y.values, torch_output, decimal=3)
+ np.testing.assert_almost_equal(np_output, torch_output, decimal=3)
def runtests(just_tests=False):
diff --git a/pysr/test/test_type_specs.py b/pysr/test/test_type_specs.py
new file mode 100644
index 000000000..e3741a0b8
--- /dev/null
+++ b/pysr/test/test_type_specs.py
@@ -0,0 +1,1797 @@
+import json
+import pickle
+import subprocess
+import sys
+import tempfile
+import unittest
+import uuid
+import warnings
+from pathlib import Path
+from unittest.mock import patch
+
+import numpy as np
+import pandas as pd
+from juliacall import JuliaError # type: ignore
+
+from pysr import PySRRegressor, TemplateExpressionSpec, TypeSpec, jl
+from pysr.expression_specs import AbstractExpressionSpec, ExpressionSpec
+from pysr.type_specs import (
+ CallableJuliaExpression,
+ compile_type_spec,
+ compile_type_spec_runtime,
+ load_type_spec_runtime,
+ object_array_1d,
+ object_array_2d,
+ prepare_type_spec_fit_data,
+ prepare_type_spec_prediction_data,
+ type_spec_to_julia_array,
+ type_spec_to_python_array,
+ validate_type_spec_configuration,
+ validate_type_spec_options,
+ validate_type_spec_runtime,
+)
+
+
+def string_spec(**overrides):
+ name = overrides.pop("name", "StringValue")
+ return TypeSpec(
+ name,
+ fields=overrides.pop("fields", {"data": "String"}),
+ sample=overrides.pop("sample", f'rng -> {name}(rand(rng, ("a", "b")))'),
+ mutate=overrides.pop(
+ "mutate",
+ f'(rng, value, temperature) -> {name}(rand(rng, ("a", "b")))',
+ ),
+ **overrides,
+ )
+
+
+def vector_spec(**overrides):
+ name = overrides.pop("name", "VectorValue")
+ return TypeSpec(
+ name,
+ fields=overrides.pop("fields", {"data": "Vector{Float64}"}),
+ sample=overrides.pop("sample", f"rng -> {name}([3.0, 4.0])"),
+ scalar_constants=overrides.pop("scalar_constants", "value -> value.data"),
+ with_scalar_constants=overrides.pop(
+ "with_scalar_constants",
+ f"(value, scalar_constants) -> {name}(collect(scalar_constants))",
+ ),
+ **overrides,
+ )
+
+
+def tiny_model(spec, *, parallelism="serial", **overrides):
+ type_name = spec.name
+ parameters = {
+ "type_spec": spec,
+ "operators": {1: [f"identity_{type_name}(x::{type_name}) = x"]},
+ "elementwise_loss": (
+ f"loss_{type_name}(x::{type_name}, y::{type_name})::Float64 = "
+ "x == y ? 0.0 : 1.0"
+ ),
+ "niterations": 1,
+ "ncycles_per_iteration": 2,
+ "populations": 1,
+ "population_size": 8,
+ "tournament_selection_n": 3,
+ "maxsize": 7,
+ "parallelism": parallelism,
+ "deterministic": parallelism == "serial",
+ "random_state": 0 if parallelism == "serial" else None,
+ "progress": False,
+ "verbosity": 0,
+ "temp_equation_file": True,
+ "should_optimize_constants": False,
+ }
+ parameters.update(overrides)
+ return PySRRegressor(**parameters)
+
+
+def identity_template():
+ return TemplateExpressionSpec(
+ combine="f(x)",
+ expressions=["f"],
+ variable_names=["x"],
+ )
+
+
+def string_data(*, constant: bool = False):
+ X = np.array([["a"], ["b"], ["a"], ["b"]], dtype=object)
+ y = np.full(len(X), "a", dtype=object) if constant else X[:, 0].copy()
+ return X, y
+
+
+class TestTypeSpecs(unittest.TestCase):
+ def test_init_invalid_value_hook(self):
+ spec = string_spec(
+ init_invalid='() -> StringValue("")',
+ is_valid="value -> !isempty(value.data)",
+ )
+ runtime = load_type_spec_runtime(compile_type_spec(spec))
+
+ self.assertTrue(
+ jl.seval(
+ "T -> begin v = SymbolicRegression.InterfaceDynamicExpressionsModule.DE.invalid_value(T); v isa T && isempty(v.data) end"
+ )(runtime.value_type)
+ )
+ X, y = string_data()
+ model = tiny_model(spec)
+ model.fit(X, y)
+ self.assertEqual(model.predict(X).tolist(), y.tolist())
+
+ def test_init_invalid_hook_rejects_wrong_type_or_valid_value(self):
+ for source, message in (
+ ('() -> ""', "init_invalid.*must return `StringValue`"),
+ (
+ '() -> StringValue("valid")',
+ "init_invalid.*must return an invalid value",
+ ),
+ ):
+ with self.subTest(source=source):
+ runtime = load_type_spec_runtime(
+ compile_type_spec(
+ string_spec(
+ init_invalid=source,
+ is_valid="value -> !isempty(value.data)",
+ )
+ )
+ )
+ with self.assertRaisesRegex(ValueError, message):
+ validate_type_spec_runtime(runtime)
+
+ def test_serial_string_type_spec(self):
+ X, y = string_data()
+ model = tiny_model(string_spec())
+ model.fit(X, y)
+
+ np.testing.assert_array_equal(model.predict(X), y)
+ self.assertTrue(
+ model._type_spec_runtime_definition_.module_name.startswith("_PySRConfig_")
+ )
+ self.assertIn("lambda_format", model.equations_.columns)
+
+ def test_named_type_is_imported_into_main(self):
+ suffix = uuid.uuid4().hex
+ type_name = f"MainVisibleValue_{suffix}"
+ runtime = load_type_spec_runtime(compile_type_spec(string_spec(name=type_name)))
+
+ self.assertTrue(bool(jl.isdefined(jl.Main, jl.Symbol(type_name))))
+ self.assertTrue(
+ bool(
+ jl.seval("(actual, expected) -> actual === expected")(
+ jl.getproperty(jl.Main, jl.Symbol(type_name)),
+ runtime.value_type,
+ )
+ )
+ )
+
+ def test_serial_vector_type_spec_with_scalar_constants(self):
+ spec = vector_spec()
+ X = np.empty((4, 1), dtype=object)
+ X[:, 0] = [
+ np.array([1.0, 2.0]),
+ np.array([3.0, 4.0]),
+ np.array([5.0, 6.0]),
+ np.array([7.0, 8.0]),
+ ]
+ model = tiny_model(spec)
+ model.fit(X, X[:, 0].copy())
+
+ self.assertEqual(
+ [list(value) for value in model.predict(X)],
+ [list(value) for value in X[:, 0]],
+ )
+ runtime = model._load_type_spec_runtime()
+ value = runtime.module._convert_value([1.0, 2.0])
+ packed, next_index = jl.seval("""
+ function (value)
+ DE = SymbolicRegression.InterfaceDynamicExpressionsModule.DE
+ buffer = fill(-99.0, 6)
+ next_index = DE.pack_scalar_constants!(buffer, 3, value)
+ return buffer, next_index
+ end
+ """)(value)
+ self.assertEqual(list(packed), [-99.0, -99.0, 1.0, 2.0, -99.0, -99.0])
+ self.assertEqual(int(next_index), 5)
+
+ def test_multithreaded_type_spec(self):
+ X, y = string_data()
+ model = tiny_model(string_spec(), parallelism="multithreading", niterations=10)
+ model.fit(X, y)
+ np.testing.assert_array_equal(model.predict(X), y)
+
+ def test_configuration_modules_isolate_colliding_operator_names(self):
+ X, y = string_data()
+ type_name = "IsolatedConfigurationValue"
+ operator_name = "colliding_identity"
+ first = tiny_model(
+ string_spec(name=type_name),
+ operators={1: [f"{operator_name}(x::{type_name}) = x"]},
+ )
+ second = tiny_model(
+ string_spec(name=type_name),
+ operators={1: [f'{operator_name}(x::{type_name}) = {type_name}("a")']},
+ )
+
+ first.fit(X, y)
+ second.fit(X, np.full(len(y), "a", dtype=object))
+
+ np.testing.assert_array_equal(first.predict(X), y)
+ np.testing.assert_array_equal(
+ second.predict(X), np.full(len(y), "a", dtype=object)
+ )
+ self.assertIsNot(
+ first._load_type_spec_runtime().configuration_module,
+ second._load_type_spec_runtime().configuration_module,
+ )
+
+ def test_preamble_struct_has_one_identity_in_hooks_and_operators(self):
+ type_name = "PreambleStructValue"
+ spec = TypeSpec(
+ type_name,
+ fields={"data": "PrivatePayload"},
+ preamble="struct PrivatePayload\n data::String\n end",
+ sample=f'rng -> {type_name}(PrivatePayload("a"))',
+ mutate="(rng, value, temperature) -> value",
+ )
+ definition = compile_type_spec_runtime(
+ spec,
+ {1: [f"x -> {type_name}(PrivatePayload(x.data.data))"]},
+ elementwise_loss="(prediction, target) -> 0.0",
+ loss_function=None,
+ loss_function_expression=None,
+ complexity_mapping=None,
+ early_stop_condition=None,
+ )
+ runtime = load_type_spec_runtime(definition)
+
+ same_object = jl.seval("(a, b) -> a === b")
+ self.assertTrue(
+ bool(
+ same_object(
+ jl.getproperty(runtime.module, jl.Symbol("PrivatePayload")),
+ jl.getproperty(
+ runtime.configuration_module, jl.Symbol("PrivatePayload")
+ ),
+ )
+ )
+ )
+ value = runtime.module._sample(runtime.module.Random.default_rng())
+ result = runtime.operator_functions[1][0](value)
+ self.assertTrue(bool(jl.isa(result.data, runtime.module.PrivatePayload)))
+
+ def test_definitions_support_fit_and_predict(self):
+ spec = TypeSpec(
+ "TransformedValue",
+ fields={"value": "Component"},
+ preamble="const Component = Float64",
+ definitions="component(value::TransformedValue) = value.value",
+ init="() -> TransformedValue(0.0)",
+ sample="rng -> TransformedValue(1.0)",
+ mutate="(rng, value, temperature) -> value",
+ )
+ X = object_array_2d([[0.0], [1.0], [2.0], [3.0]])
+ y = object_array_1d([2.0, 4.0, 6.0, 8.0])
+ model = tiny_model(
+ spec,
+ operators={
+ 1: [
+ """
+ transform(value::TransformedValue)::TransformedValue =
+ TransformedValue(2 * (component(value) + 1))
+ """
+ ]
+ },
+ elementwise_loss="""
+ transformed_loss(a::TransformedValue, b::TransformedValue)::Float64 =
+ abs2(component(a) - component(b))
+ """,
+ guesses=["transform(x0)"],
+ )
+ model.fit(X, y)
+ np.testing.assert_array_equal(model.predict(X), y)
+
+ def test_definitions_are_private_to_their_specification(self):
+ type_name = "DefinitionIdentityValue"
+ for tag in ("first", "second"):
+ with self.subTest(tag=tag):
+ runtime = load_type_spec_runtime(
+ compile_type_spec(
+ string_spec(
+ name=type_name,
+ definitions=f'tag(::{type_name}) = "{tag}"',
+ )
+ )
+ )
+ value = runtime.module._sample(runtime.module.Random.default_rng())
+ self.assertEqual(str(runtime.module.tag(value)), tag)
+
+ def test_repeated_runtime_load_does_not_mutate_child_module(self):
+ spec = string_spec(name="RepeatedLoadValue")
+ definition = compile_type_spec_runtime(
+ spec,
+ {1: ["x -> x"]},
+ elementwise_loss="(prediction, target) -> 0.0",
+ loss_function=None,
+ loss_function_expression=None,
+ complexity_mapping=None,
+ early_stop_condition=None,
+ expression_spec="x -> x",
+ expression_spec_function_selector="spec -> spec",
+ )
+ first = load_type_spec_runtime(definition)
+ names_before = tuple(
+ str(name)
+ for name in jl.names(first.configuration_module, all=True, imported=True)
+ )
+
+ second = load_type_spec_runtime(first.definition)
+ names_after = tuple(
+ str(name)
+ for name in jl.names(second.configuration_module, all=True, imported=True)
+ )
+
+ self.assertEqual(names_after, names_before)
+ self.assertTrue(
+ bool(
+ jl.seval("(a, b) -> a === b")(
+ first.expression_spec,
+ second.expression_spec,
+ )
+ )
+ )
+
+ def test_failed_configuration_source_remains_isolated(self):
+ suffix = uuid.uuid4().hex
+ type_name = f"FailedConfigurationValue_{suffix}"
+ operator_name = f"isolated_operator_{suffix}"
+ X, y = string_data()
+ fitted = tiny_model(
+ string_spec(name=type_name),
+ operators={1: [f"{operator_name}(x::{type_name}) = x"]},
+ )
+ fitted.fit(X, y)
+ broken = tiny_model(
+ string_spec(name=type_name),
+ operators={
+ 1: [
+ f"""
+ begin
+ {operator_name}(x::{type_name}) = {type_name}("a")
+ error("isolated source failure")
+ end
+ """
+ ]
+ },
+ )
+
+ with self.assertRaisesRegex(JuliaError, "isolated source failure"):
+ broken.fit(X, y)
+
+ np.testing.assert_array_equal(fitted.predict(X), y)
+ self.assertFalse(bool(jl.isdefined(jl.Main, jl.Symbol(operator_name))))
+
+ def test_failed_runtime_install_can_retry_same_definition(self):
+ suffix = uuid.uuid4().hex
+ flag = f"_pysr_retry_install_{suffix}"
+ type_name = f"RetryInstallValue_{suffix}"
+ operator_name = f"retry_identity_{suffix}"
+ jl.seval(f"global {flag} = Ref(true)")
+ definition = compile_type_spec_runtime(
+ string_spec(name=type_name),
+ {1: [f"{operator_name}(x::{type_name})::{type_name} = x"]},
+ elementwise_loss=f"(x::{type_name}, y::{type_name}) -> 0.0",
+ loss_function=None,
+ loss_function_expression=None,
+ complexity_mapping=None,
+ early_stop_condition=f"""
+ begin
+ Main.{flag}[] && error("retry source failure")
+ (loss, complexity) -> false
+ end
+ """,
+ )
+
+ with self.assertRaisesRegex(JuliaError, "retry source failure"):
+ load_type_spec_runtime(definition)
+
+ jl.seval(f"{flag}[] = false")
+ runtime = load_type_spec_runtime(definition)
+ self.assertEqual(runtime.operator_names, {1: (operator_name,)})
+ self.assertIsNotNone(runtime.early_stop_condition)
+
+ def test_failed_type_install_can_retry_same_definition(self):
+ suffix = uuid.uuid4().hex
+ flag = f"_pysr_retry_type_install_{suffix}"
+ type_name = f"RetryTypeValue_{suffix}"
+ operator_name = f"retry_type_identity_{suffix}"
+ jl.seval(f"global {flag} = Ref(true)")
+ definition = compile_type_spec_runtime(
+ string_spec(
+ name=type_name,
+ preamble=f'Main.{flag}[] && error("retry type failure")',
+ ),
+ {1: [f"{operator_name}(x::{type_name})::{type_name} = x"]},
+ elementwise_loss=f"(x::{type_name}, y::{type_name}) -> 0.0",
+ loss_function=None,
+ loss_function_expression=None,
+ complexity_mapping=None,
+ early_stop_condition=None,
+ )
+
+ with self.assertRaisesRegex(JuliaError, "retry type failure"):
+ load_type_spec_runtime(definition)
+
+ jl.seval(f"{flag}[] = false")
+ runtime = load_type_spec_runtime(definition)
+ self.assertEqual(runtime.operator_names, {1: (operator_name,)})
+ self.assertTrue(
+ bool(jl.isdefined(runtime.module, jl.Symbol(type_name))),
+ )
+
+ def test_equivalent_runtime_definition_is_evaluated_once(self):
+ suffix = uuid.uuid4().hex
+ counter = f"_pysr_definition_count_{suffix}"
+ type_name = f"ReusedValue_{suffix}"
+ operator = f"reused_identity_{suffix}"
+ jl.seval(f"global {counter} = 0")
+ source = f"""
+ begin
+ Main.{counter} += 1
+ {operator}(x::{type_name}) = x
+ end
+ """
+ first = tiny_model(string_spec(name=type_name), operators={1: [source]})
+ second = tiny_model(string_spec(name=type_name), operators={1: [source]})
+ X, y = string_data()
+
+ first.fit(X, y)
+ second.fit(X, y)
+ restored = pickle.loads(pickle.dumps(second))
+ _ = restored.julia_state_
+ _ = restored.julia_options_
+
+ self.assertEqual(int(jl.seval(counter)), 1)
+ self.assertTrue(
+ bool(
+ jl.seval("(a, b) -> a === b")(
+ first._load_type_spec_runtime().configuration_module,
+ second._load_type_spec_runtime().configuration_module,
+ )
+ )
+ )
+
+ def test_restored_model_rebuilds_callable_columns(self):
+ X, y = string_data()
+ model = tiny_model(string_spec())
+ model.fit(X, y)
+
+ restored = pickle.loads(pickle.dumps(model))
+ self.assertNotIn("lambda_format", restored.equations_.columns)
+
+ best = restored.get_best()
+ self.assertIn("lambda_format", restored.equations_.columns)
+ np.testing.assert_array_equal(best["lambda_format"](X), model.predict(X))
+
+ def test_failed_warm_start_restores_fitted_metadata_and_checkpoint(self):
+ X, y = string_data()
+ with tempfile.TemporaryDirectory() as directory:
+ model = tiny_model(
+ string_spec(),
+ temp_equation_file=False,
+ output_directory=directory,
+ run_id="warm-start-rollback",
+ )
+ model.fit(X, y)
+ definition = model._type_spec_runtime_definition_
+ state = model.julia_state_stream_.copy()
+ equations = model.equations_.copy(deep=True)
+ feature_names = model.feature_names_in_.copy()
+ checkpoint = Path(model.get_pkl_filename())
+ checkpoint_contents = checkpoint.read_bytes()
+
+ model.set_params(
+ warm_start=True,
+ complexity_mapping="expression -> 1",
+ )
+ with self.assertRaisesRegex(
+ ValueError, "Cannot warm-start after changing.*configuration"
+ ):
+ model.fit(X, y, variable_names=["candidate"])
+
+ self.assertEqual(model._type_spec_runtime_definition_, definition)
+ np.testing.assert_array_equal(model.julia_state_stream_, state)
+ pd.testing.assert_frame_equal(model.equations_, equations)
+ np.testing.assert_array_equal(model.feature_names_in_, feature_names)
+ self.assertEqual(checkpoint.read_bytes(), checkpoint_contents)
+
+ def fail_after_mutation(*args, **kwargs):
+ model.equations_.iloc[0, 0] += 100
+ model.feature_names_in_[0] = "mutated"
+ raise RuntimeError("search failed")
+
+ model.set_params(complexity_mapping=None)
+ with (
+ patch.object(model, "_run", side_effect=fail_after_mutation),
+ self.assertRaisesRegex(RuntimeError, "search failed"),
+ ):
+ model.fit(X, y, variable_names=["candidate"])
+
+ self.assertEqual(model._type_spec_runtime_definition_, definition)
+ np.testing.assert_array_equal(model.julia_state_stream_, state)
+ pd.testing.assert_frame_equal(model.equations_, equations)
+ np.testing.assert_array_equal(model.feature_names_in_, feature_names)
+ self.assertEqual(checkpoint.read_bytes(), checkpoint_contents)
+
+ def test_warm_start_rejects_loss_type_change(self):
+ X, y = string_data()
+ type_name = "WarmStartLossTypeValue"
+ objective = """
+ function warm_start_objective(
+ tree, dataset::Dataset, options, idx=nothing
+ )::Float64
+ _, complete = eval_tree_array(tree, dataset.X, options)
+ return complete ? 0.0 : Inf
+ end
+ """
+ model = tiny_model(
+ string_spec(name=type_name, loss_type="Float64"),
+ elementwise_loss=None,
+ loss_function=objective,
+ )
+ model.fit(X, y)
+ model.set_params(
+ warm_start=True,
+ type_spec=string_spec(name=type_name, loss_type="Float32"),
+ )
+
+ with self.assertRaisesRegex(
+ ValueError, "Cannot warm-start after changing.*configuration"
+ ):
+ model.fit(X, y)
+
+ def test_pre_search_checkpoint_can_resume_as_a_fresh_fit(self):
+ X, y = string_data()
+ with tempfile.TemporaryDirectory() as directory:
+ model = tiny_model(
+ string_spec(name="PreSearchCheckpointValue"),
+ temp_equation_file=False,
+ output_directory=directory,
+ run_id="pre-search-checkpoint",
+ )
+ with (
+ patch.object(
+ PySRRegressor,
+ "_run",
+ side_effect=RuntimeError("search failed"),
+ ),
+ self.assertRaisesRegex(RuntimeError, "search failed"),
+ ):
+ model.fit(X, y)
+
+ checkpoint = Path(model.get_pkl_filename())
+ with checkpoint.open("rb") as checkpoint_file:
+ restored = pickle.load(checkpoint_file)
+ self.assertIsNone(restored.equations_)
+ self.assertFalse(hasattr(restored, "_type_spec_runtime_definition_"))
+
+ restored.set_params(
+ temp_equation_file=True,
+ output_directory=None,
+ run_id=None,
+ )
+ restored.fit(X, y)
+ np.testing.assert_array_equal(restored.predict(X), y)
+
+ def test_callable_supervised_loss_is_supported(self):
+ type_name = "CallableLossValue"
+ loss_name = "CallableTypeSpecLoss"
+ X, y = string_data()
+ model = tiny_model(
+ string_spec(name=type_name),
+ elementwise_loss=f"""
+ begin
+ struct {loss_name} <: SymbolicRegression.SupervisedLoss end
+ (::{loss_name})(
+ prediction::{type_name}, target::{type_name}
+ )::Float64 = prediction == target ? 0.0 : 1.0
+ {loss_name}()
+ end
+ """,
+ )
+
+ model.fit(X, y)
+ np.testing.assert_array_equal(model.predict(X), y)
+
+ def test_numeric_early_stop_condition_remains_numeric(self):
+ model = tiny_model(
+ string_spec(name="NumericEarlyStopValue"),
+ early_stop_condition=1e-4,
+ )
+ definition = model._compile_type_spec_runtime(model._operators_from_params())
+ runtime = load_type_spec_runtime(definition)
+
+ self.assertEqual(float(runtime.early_stop_condition), 1e-4)
+
+ def test_pickle_restores_runtime_before_deserializing_in_fresh_process(self):
+ X, y = string_data()
+ with tempfile.TemporaryDirectory() as directory:
+ model = tiny_model(
+ string_spec(name="CheckpointStringValue"),
+ complexity_mapping="checkpoint_complexity(expression) = 1",
+ early_stop_condition=(
+ "checkpoint_early_stop(loss, complexity) = false"
+ ),
+ temp_equation_file=False,
+ output_directory=directory,
+ run_id="typespec-checkpoint",
+ delete_tempfiles=False,
+ )
+ model.fit(X, y)
+ type_module_name = (
+ model._type_spec_runtime_definition_.type_spec.module_name
+ )
+ run_directory = Path(directory) / "typespec-checkpoint"
+ module_name = model._type_spec_runtime_definition_.module_name
+ code = f"""
+import json
+import numpy as np
+from pysr import PySRRegressor, jl
+type_module_name = {type_module_name!r}
+module_name = {module_name!r}
+assert not bool(jl.isdefined(jl.Main, jl.Symbol(type_module_name)))
+model = PySRRegressor.from_file(run_directory={str(run_directory)!r})
+X = np.array([["a"], ["b"], ["a"], ["b"]], dtype=object)
+runtime = model._load_type_spec_runtime()
+print(json.dumps({{
+ "prediction": model.predict(X).tolist(),
+ "loaded": bool(jl.isdefined(runtime.module, jl.Symbol(module_name))),
+ "complexity": str(jl.Base.nameof(runtime.complexity_mapping)),
+ "early_stop": str(jl.Base.nameof(runtime.early_stop_condition)),
+}}))
+"""
+ result = subprocess.run(
+ [sys.executable, "-c", code],
+ check=False,
+ capture_output=True,
+ text=True,
+ )
+
+ self.assertEqual(result.returncode, 0, result.stderr)
+ payload = json.loads(result.stdout.strip().splitlines()[-1])
+ self.assertEqual(
+ payload,
+ {
+ "prediction": y.tolist(),
+ "loaded": True,
+ "complexity": "checkpoint_complexity",
+ "early_stop": "checkpoint_early_stop",
+ },
+ )
+
+ def test_unsupported_configurations_fail_before_runtime_loading(self):
+ X, y = string_data()
+ cases = (
+ (tiny_model(string_spec(), turbo=True), ValueError, "turbo"),
+ (tiny_model(string_spec(), bumper=True), ValueError, "bumper"),
+ (
+ tiny_model(string_spec(), autodiff_backend="Zygote"),
+ ValueError,
+ "autodiff_backend",
+ ),
+ (
+ tiny_model(string_spec(), output_jax_format=True),
+ ValueError,
+ "output_jax_format",
+ ),
+ (
+ tiny_model(string_spec(), output_torch_format=True),
+ ValueError,
+ "output_torch_format",
+ ),
+ (
+ tiny_model(string_spec(), extra_sympy_mappings={"f": lambda x: x}),
+ ValueError,
+ "extra_sympy_mappings",
+ ),
+ (
+ tiny_model(
+ string_spec(), extra_jax_mappings={lambda x: x: "jnp.asarray"}
+ ),
+ ValueError,
+ "extra_jax_mappings",
+ ),
+ (
+ tiny_model(
+ string_spec(), extra_torch_mappings={lambda x: x: lambda x: x}
+ ),
+ ValueError,
+ "extra_torch_mappings",
+ ),
+ (
+ tiny_model(string_spec(), binary_operators=["+"]),
+ ValueError,
+ "binary_operators",
+ ),
+ )
+ for model, error, message in cases:
+ with self.subTest(message=message):
+ with self.assertRaisesRegex(error, message):
+ model.fit(X, y)
+ self.assertFalse(hasattr(model, "_type_spec_runtime_definition_"))
+
+ def test_custom_expression_spec_requires_runtime_source(self):
+ class MissingSourceSpec(ExpressionSpec):
+ pass
+
+ X, y = string_data()
+ model = tiny_model(
+ string_spec(name="MissingExpressionSourceValue"),
+ expression_spec=MissingSourceSpec(),
+ )
+
+ with self.assertRaisesRegex(ValueError, "does not provide Julia source"):
+ model.fit(X, y)
+ self.assertFalse(hasattr(model, "_type_spec_runtime_definition_"))
+
+ def test_operator_names_drive_constraint_keys(self):
+ type_name = "NamedOperatorValue"
+ operator_name = "named_identity"
+ X, y = string_data()
+ model = tiny_model(
+ string_spec(name=type_name),
+ operators={1: [f"{operator_name}(x::{type_name})::{type_name} = x"]},
+ constraints={operator_name: 1},
+ )
+
+ model.fit(X, y)
+
+ runtime = model._load_type_spec_runtime()
+ self.assertEqual(runtime.operator_names, {1: (operator_name,)})
+
+ def test_template_source_hook_runs_once_per_fit(self):
+ sources = []
+
+ class SourceExpressionSpec(AbstractExpressionSpec):
+ @property
+ def supports_type_spec(self):
+ return True
+
+ def julia_expression_spec(self):
+ return ExpressionSpec().julia_expression_spec()
+
+ def create_exports(self, *args, **kwargs):
+ return ExpressionSpec().create_exports(*args, **kwargs)
+
+ def _julia_expression_spec_source(self, *, prototype=None):
+ sources.append(prototype)
+ return f"""
+ @template_spec(
+ expressions=(f,), prototype={prototype}
+ ) do x
+ f(x)
+ end
+ """
+
+ def _julia_expression_spec_function_selector(self):
+ return "spec -> spec.structure.combine"
+
+ X, y = string_data()
+ model = tiny_model(
+ string_spec(name="SingleTemplateSourceValue"),
+ expression_spec=SourceExpressionSpec(),
+ )
+ model.fit(X, y)
+
+ self.assertEqual(
+ sources,
+ ["SymbolicRegression.init_value(SingleTemplateSourceValue)"],
+ )
+ np.testing.assert_array_equal(model.predict(X), y)
+
+ def test_template_type_spec_parallelism_matrix(self):
+ type_name = "TemplateParallelismValue"
+ for parallelism in ("serial", "multithreading", "multiprocessing"):
+ X, y = string_data(constant=True)
+ model = tiny_model(
+ string_spec(
+ name=type_name,
+ sample=f'rng -> {type_name}("a")',
+ mutate="(rng, value, temperature) -> value",
+ ),
+ expression_spec=identity_template(),
+ parallelism=parallelism,
+ procs=2 if parallelism == "multiprocessing" else None,
+ niterations=100,
+ early_stop_condition=1e-8,
+ )
+
+ with self.subTest(parallelism=parallelism):
+ model.fit(X, y)
+ np.testing.assert_array_equal(model.predict(X), y)
+
+ def test_same_runtime_warm_start_continues(self):
+ X, y = string_data()
+ model = tiny_model(
+ string_spec(),
+ expression_spec=identity_template(),
+ )
+ model.fit(X, y)
+ fingerprint = model._type_spec_runtime_definition_.fingerprint
+ model.set_params(warm_start=True)
+
+ model.fit(X, y)
+
+ self.assertEqual(
+ model._type_spec_runtime_definition_.fingerprint,
+ fingerprint,
+ )
+ np.testing.assert_array_equal(model.predict(X), y)
+
+ def test_malformed_hooks_fail_with_hook_names(self):
+ cases = (
+ (
+ vector_spec(
+ name="BadScalarConstantsValue",
+ scalar_constants="value -> 1.0",
+ ),
+ "scalar_constants.*AbstractVector",
+ ),
+ (
+ string_spec(
+ name="BadPredicateValue",
+ is_valid="value -> 1",
+ ),
+ "is_valid.*Bool",
+ ),
+ (
+ string_spec(name="BadInitValue", init='() -> ""'),
+ "init.*BadInitValue",
+ ),
+ (
+ vector_spec(
+ name="BadReconstructionValue",
+ with_scalar_constants="(value, constants) -> constants",
+ ),
+ "with_scalar_constants.*BadReconstructionValue",
+ ),
+ (
+ vector_spec(
+ name="BadScalarConstantElementValue",
+ scalar_constants="value -> [1, 2]",
+ with_scalar_constants=(
+ "(value, constants) -> "
+ "BadScalarConstantElementValue(collect(constants))"
+ ),
+ ),
+ "scalar_constants.*concrete `AbstractFloat`",
+ ),
+ (
+ vector_spec(
+ name="NonRoundTripValue",
+ with_scalar_constants=(
+ "(value, constants) -> "
+ "NonRoundTripValue(reverse(collect(constants)))"
+ ),
+ ),
+ "with_scalar_constants.*preserve",
+ ),
+ (
+ vector_spec(name="BadStringValue", string="value -> 1"),
+ "string.*AbstractString",
+ ),
+ (
+ string_spec(
+ name="BadSampleArityValue",
+ sample='() -> BadSampleArityValue("a")',
+ ),
+ "sample",
+ ),
+ (
+ string_spec(
+ name="RejectedSampleValue",
+ is_valid="value -> false",
+ ),
+ "sample.*invalid value",
+ ),
+ )
+ for spec, message in cases:
+ with self.subTest(spec=spec.name):
+ runtime = load_type_spec_runtime(compile_type_spec(spec))
+ with self.assertRaisesRegex(ValueError, message):
+ validate_type_spec_runtime(runtime)
+
+ def test_first_fit_validates_type_spec_hooks(self):
+ X, y = string_data()
+ model = tiny_model(string_spec(name="FitValidationValue", init='() -> ""'))
+
+ with self.assertRaisesRegex(ValueError, "init.*FitValidationValue"):
+ model.fit(X, y)
+ self.assertFalse(hasattr(model, "_type_spec_runtime_definition_"))
+
+ def test_operator_and_elementwise_loss_are_invoked_during_validation(self):
+ type_name = "RuntimeInvocationValue"
+ cases = {
+ "operator": {
+ "operator": (
+ f"function exploding_operator(x::{type_name})::{type_name}\n"
+ ' x.data == "b" && error("operator was invoked")\n'
+ " return x\nend"
+ ),
+ "loss": (
+ f"identity_loss(x::{type_name}, y::{type_name})::Float64 = 0.0"
+ ),
+ "message": "operator was invoked",
+ },
+ "elementwise loss": {
+ "operator": (f"identity_runtime(x::{type_name})::{type_name} = x"),
+ "loss": (
+ f"function exploding_loss(x::{type_name}, y::{type_name})::Float64\n"
+ ' x.data == y.data && error("loss was invoked")\n'
+ " return 0.0\nend"
+ ),
+ "message": "loss was invoked",
+ },
+ }
+ for name, case in cases.items():
+ with self.subTest(case=name):
+ definition = compile_type_spec_runtime(
+ string_spec(name=type_name),
+ {1: [case["operator"]]},
+ elementwise_loss=case["loss"],
+ loss_function=None,
+ loss_function_expression=None,
+ complexity_mapping=None,
+ early_stop_condition=None,
+ )
+ runtime = load_type_spec_runtime(definition)
+ with self.assertRaisesRegex(ValueError, case["message"]):
+ validate_type_spec_runtime(runtime)
+
+ def test_module_loading_has_no_behavioral_hook_side_effects(self):
+ suffix = uuid.uuid4().hex
+ init_count = f"_pysr_init_count_{suffix}"
+ mutate_count = f"_pysr_mutate_count_{suffix}"
+ type_name = f"SideEffectValue_{suffix}"
+ jl.seval(f"global {init_count} = 0; global {mutate_count} = 0")
+ spec = string_spec(
+ name=type_name,
+ init=(f"() -> begin Main.{init_count} += 1; " f'{type_name}("a") end'),
+ mutate=(
+ f"(rng, value, temperature) -> begin Main.{mutate_count} += 1; "
+ "value end"
+ ),
+ )
+
+ load_type_spec_runtime(compile_type_spec(spec))
+ self.assertEqual(int(jl.seval(init_count)), 0)
+ self.assertEqual(int(jl.seval(mutate_count)), 0)
+
+ definition = compile_type_spec_runtime(
+ spec,
+ {1: [f"side_effect_identity(x::{type_name}) = x"]},
+ elementwise_loss=(
+ f"side_effect_loss(x::{type_name}, y::{type_name})::Float64 = 0.0"
+ ),
+ loss_function=None,
+ loss_function_expression=None,
+ complexity_mapping=None,
+ early_stop_condition=None,
+ )
+ runtime = load_type_spec_runtime(definition)
+ validate_type_spec_options(
+ runtime,
+ runtime.operator_functions,
+ runtime.elementwise_loss,
+ )
+ self.assertEqual(int(jl.seval(init_count)), 0)
+ self.assertEqual(int(jl.seval(mutate_count)), 0)
+ validate_type_spec_runtime(runtime)
+ self.assertEqual(int(jl.seval(init_count)), 1)
+ self.assertEqual(int(jl.seval(mutate_count)), 1)
+
+ def test_object_array_shapes_and_logical_conversion(self):
+ one_dimensional = object_array_1d(value for value in ([1, 2], [3, 4]))
+ self.assertEqual(one_dimensional.shape, (2,))
+ self.assertEqual(one_dimensional.tolist(), [[1, 2], [3, 4]])
+ two_dimensional = object_array_2d(
+ row for row in (([1, 2], [3, 4]), ([5, 6], [7, 8]))
+ )
+ self.assertEqual(two_dimensional.shape, (2, 2))
+ with self.assertRaisesRegex(ValueError, "2D array"):
+ object_array_2d(["a", "b"])
+ with self.assertRaisesRegex(ValueError, "same number"):
+ object_array_2d([[1], [2, 3]])
+
+ spec = TypeSpec(
+ "PairValue",
+ fields={"number": "Float64", "label": "String"},
+ sample='rng -> PairValue(0.0, "")',
+ mutate="(rng, value, temperature) -> value",
+ )
+ runtime = load_type_spec_runtime(compile_type_spec(spec))
+ values = object_array_1d([(1.0, "one"), (2.0, "two")])
+ converted = type_spec_to_julia_array(runtime, values)
+ self.assertEqual(
+ type_spec_to_python_array(runtime, converted).tolist(), values.tolist()
+ )
+
+ def test_derived_scalar_constant_packing_honors_offsets(self):
+ runtime = load_type_spec_runtime(
+ compile_type_spec(vector_spec(name="OffsetVectorValue"))
+ )
+ first = runtime.module._convert_value([1.0, 2.0])
+ second = runtime.module._convert_value([3.0, 4.0])
+ packed, next_index, unpacked = jl.seval("""
+ function (first, second)
+ DE = SymbolicRegression.InterfaceDynamicExpressionsModule.DE
+ packed = fill(-99.0, 8)
+ idx = DE.pack_scalar_constants!(packed, 3, first)
+ idx = DE.pack_scalar_constants!(packed, idx, second)
+ idx, rebuilt_first =
+ DE.unpack_scalar_constants(packed, 3, first)
+ idx, rebuilt_second =
+ DE.unpack_scalar_constants(packed, idx, second)
+ return packed, idx, (rebuilt_first, rebuilt_second)
+ end
+ """)(first, second)
+ self.assertEqual(
+ list(packed),
+ [-99.0, -99.0, 1.0, 2.0, 3.0, 4.0, -99.0, -99.0],
+ )
+ self.assertEqual(int(next_index), 7)
+ self.assertEqual(
+ [list(value.data) for value in unpacked],
+ [[1.0, 2.0], [3.0, 4.0]],
+ )
+
+ def test_union_payload_scalar_constant_hooks(self):
+ spec = TypeSpec(
+ "TensorValue",
+ fields={"data": "Union{Float64, Vector{Float64}, Matrix{Float64}}"},
+ sample="rng -> TensorValue(randn(rng, 2, 2))",
+ scalar_constants="""
+ function scalar_constants(value)
+ return value.data isa Float64 ? [value.data] : vec(value.data)
+ end
+ """,
+ with_scalar_constants="""
+ function with_scalar_constants(value, scalar_constants)
+ data = value.data isa Float64 ? scalar_constants[1] :
+ reshape(collect(scalar_constants), size(value.data))
+ return TensorValue(data)
+ end
+ """,
+ )
+ runtime = load_type_spec_runtime(compile_type_spec(spec))
+ value = runtime.module._convert_value(np.eye(2))
+ self.assertEqual(
+ list(runtime.module._scalar_constants(value)),
+ [1.0, 0.0, 0.0, 1.0],
+ )
+
+ def test_immutable_scalar_constant_view_validates(self):
+ spec = vector_spec(
+ name="ImmutableScalarConstantsValue",
+ scalar_constants=(
+ "value -> range(value.data[1], value.data[end]; "
+ "length=length(value.data))"
+ ),
+ with_scalar_constants=(
+ "(value, scalar_constants) -> "
+ "ImmutableScalarConstantsValue(collect(scalar_constants))"
+ ),
+ )
+ model = tiny_model(spec)
+ runtime = model._load_type_spec_runtime(
+ for_fit=True,
+ operators=model._operators_from_params(),
+ )
+ self.assertEqual(runtime.spec.name, "ImmutableScalarConstantsValue")
+
+ def test_variable_length_vector_constants(self):
+ spec = TypeSpec(
+ "VariableVector",
+ fields={"data": "Vector{Float64}"},
+ sample="rng -> VariableVector(randn(rng, rand(rng, 1:5)))",
+ mutate="""
+ function mutate_vector(rng, value, temperature)
+ if rand(rng) < 0.2
+ return VariableVector(randn(rng, rand(rng, 1:5)))
+ end
+ data = copy(value.data)
+ data[rand(rng, eachindex(data))] += temperature * randn(rng)
+ return VariableVector(data)
+ end
+ """,
+ scalar_constants="value -> value.data",
+ with_scalar_constants=(
+ "(value, scalar_constants) -> "
+ "VariableVector(collect(scalar_constants))"
+ ),
+ )
+ rng = np.random.default_rng(0)
+ values = [rng.normal(size=2) for _ in range(64)]
+ X = pd.DataFrame({"x": values})
+ prefix = np.array([1.5])
+ suffix = np.array([-0.2, -3.0, 0.1])
+ y = np.empty(len(values), dtype=object)
+ y[:] = [np.concatenate((prefix, value, suffix)) for value in values]
+ model = tiny_model(
+ spec,
+ expression_spec=identity_template(),
+ operators={
+ 2: ["concat_vectors(a, b) = " "VariableVector(vcat(a.data, b.data))"]
+ },
+ elementwise_loss="""
+ function vector_loss(a, b)::Float64
+ return length(a.data) == length(b.data) ?
+ sum(abs2, a.data - b.data) : 1.0e6
+ end
+ """,
+ niterations=60,
+ ncycles_per_iteration=100,
+ populations=4,
+ population_size=50,
+ tournament_selection_n=10,
+ maxsize=7,
+ early_stop_condition=(
+ "(loss, complexity) -> loss < 1.0e-8 && complexity == 5"
+ ),
+ should_optimize_constants=True,
+ )
+
+ model.fit(X, y)
+
+ exact = model.equations_.query("complexity == 5").sort_values("loss").iloc[0]
+ self.assertLess(exact.loss, 1.0e-8)
+ constants = jl.seval("""
+ function (expression)
+ DE = SymbolicRegression.InterfaceDynamicExpressionsModule.DE
+ tree = DE.get_tree(expression)
+ nodes = DE.filter_map(
+ node -> node.degree == 0 && node.constant,
+ identity,
+ tree,
+ typeof(tree),
+ )
+ return map(node -> node.val.data, nodes)
+ end
+ """)(exact.julia_expression)
+ self.assertEqual([len(value) for value in constants], [1, 3])
+ np.testing.assert_allclose(constants[0], prefix, atol=1.0e-6)
+ np.testing.assert_allclose(constants[1], suffix, atol=1.0e-6)
+
+ def test_multi_field_fit_predicts_tuples(self):
+ spec = TypeSpec(
+ "FitPairValue",
+ fields={"number": "Float64", "label": "String"},
+ sample='rng -> FitPairValue(0.0, "")',
+ mutate="(rng, value, temperature) -> value",
+ )
+ pairs = [(1.0, "one"), (2.0, "two"), (3.0, "three"), (4.0, "four")]
+ X = pd.DataFrame({"x": pairs})
+ y = pd.Series(pairs, dtype=object)
+ model = tiny_model(spec)
+
+ model.fit(X, y)
+
+ self.assertEqual(model.predict(X).tolist(), pairs)
+
+ def test_invalid_expression_evaluation_reports_the_rejected_value(self):
+ type_name = "InvalidEvaluationValue"
+ spec = TypeSpec(
+ type_name,
+ fields={"data": "Float64"},
+ sample=f"rng -> {type_name}(1.0)",
+ scalar_constants="value -> (value.data,)",
+ with_scalar_constants=f"(value, constants) -> {type_name}(constants[1])",
+ )
+ definition = compile_type_spec_runtime(
+ spec,
+ {1: [f"invalid_operator(x::{type_name})::{type_name} = {type_name}(Inf)"]},
+ elementwise_loss="(prediction, target) -> 0.0",
+ loss_function=None,
+ loss_function_expression=None,
+ complexity_mapping=None,
+ early_stop_condition=None,
+ )
+ runtime = load_type_spec_runtime(definition)
+ expression = jl.seval("""
+ (value_type, operator) -> begin
+ operators = OperatorEnum(1 => (operator,))
+ node = Node{value_type}(op=1, l=Node{value_type}(feature=1))
+ Expression(node; operators, variable_names=["x1"])
+ end
+ """)(runtime.value_type, runtime.operator_functions[1][0])
+ X = object_array_2d([[1.0], [2.0]])
+
+ with self.assertRaisesRegex(ValueError, "`is_valid` rejected"):
+ CallableJuliaExpression(expression, runtime)(X)
+
+ def test_type_spec_variable_name_count_is_validated(self):
+ X, y = string_data()
+ model = tiny_model(string_spec(name="WrongVariableCountValue"))
+
+ with self.assertRaisesRegex(
+ ValueError, "variable_names.*one name per TypeSpec feature"
+ ):
+ model.fit(X, y, variable_names=["first", "extra"])
+
+ def test_rejects_empty_feature_axis(self):
+ model = tiny_model(string_spec())
+
+ with self.assertRaisesRegex(ValueError, "at least one feature"):
+ model.fit(np.empty((2, 0), dtype=object), np.array(["a", "b"]))
+
+ def test_csv_only_loading_is_rejected(self):
+ with tempfile.TemporaryDirectory() as directory:
+ Path(directory, "hall_of_fame.csv").write_text(
+ "Complexity,Loss,Equation\n1,0.0,x0\n",
+ encoding="utf-8",
+ )
+ with self.assertRaisesRegex(ValueError, "checkpoint.pkl"):
+ PySRRegressor.from_file(
+ run_directory=directory,
+ type_spec=string_spec(),
+ operators={1: ["identity_value(x::StringValue) = x"]},
+ n_features_in=1,
+ )
+
+ def test_template_custom_combiner_infers_num_features(self):
+ type_name = "TemplateVectorValue"
+ model = tiny_model(
+ vector_spec(name=type_name),
+ expression_spec=TemplateExpressionSpec(
+ combine="add_vectors(f(x1), g(x2))",
+ expressions=["f", "g"],
+ variable_names=["x1", "x2"],
+ ),
+ operators={
+ 2: [
+ f"""
+ add_vectors(a::{type_name}, b::{type_name}) =
+ {type_name}(a.data + b.data)
+ add_vectors(a::ValidVector, b::ValidVector) =
+ ValidVector(map(add_vectors, a.x, b.x), a.valid && b.valid)
+ """
+ ]
+ },
+ elementwise_loss=(
+ f"vector_loss(a::{type_name}, b::{type_name})::Float64 = "
+ "sum(abs2, a.data - b.data)"
+ ),
+ )
+ X = np.empty((4, 2), dtype=object)
+ X[:, 0] = [np.array([1.0, 2.0])] * 4
+ X[:, 1] = [np.array([3.0, 4.0])] * 4
+ y = np.empty(4, dtype=object)
+ y[:] = [np.array([4.0, 6.0])] * 4
+
+ model.fit(X, y, variable_names=["x1", "x2"])
+
+ structure = model.julia_options_.expression_options.structure
+ self.assertEqual(int(structure.num_features.f), 1)
+ self.assertEqual(int(structure.num_features.g), 1)
+
+ def test_template_type_spec_parameters(self):
+ type_name = "TemplateParameterValue"
+ X, y = string_data(constant=True)
+ expression_spec = TemplateExpressionSpec(
+ combine="choose_parameter(p[1], f(x))",
+ expressions=["f"],
+ variable_names=["x"],
+ parameters={"p": 1},
+ )
+ model = tiny_model(
+ string_spec(
+ name=type_name,
+ sample=f'rng -> {type_name}("a")',
+ mutate="(rng, value, temperature) -> value",
+ ),
+ expression_spec=expression_spec,
+ operators={
+ 1: [f"identity_{type_name}(x::{type_name}) = x"],
+ 2: [
+ f"""
+ choose_parameter(a::{type_name}, b::{type_name}) = a
+ choose_parameter(a::{type_name}, b::ValidVector) =
+ ValidVector(map(_ -> a, b.x), b.valid)
+ """
+ ],
+ },
+ parallelism="multiprocessing",
+ procs=2,
+ )
+
+ model.fit(X, y)
+ fingerprint = model._type_spec_runtime_definition_.fingerprint
+ model.set_params(warm_start=True)
+
+ model.fit(X, y)
+
+ self.assertEqual(model._type_spec_runtime_definition_.fingerprint, fingerprint)
+ np.testing.assert_array_equal(model.predict(X), y)
+
+ def test_template_type_spec_parameters_are_mutated(self):
+ type_name = "MutatedParameterValue"
+ X = np.array([["a"], ["b"], ["a"], ["b"]], dtype=object)
+ y = np.full(len(X), "b", dtype=object)
+ model = tiny_model(
+ string_spec(
+ name=type_name,
+ sample=f'rng -> {type_name}("a")',
+ mutate=f'(rng, value, temperature) -> {type_name}("b")',
+ ),
+ expression_spec=TemplateExpressionSpec(
+ combine="choose_parameter(p[1], f(x))",
+ expressions=["f"],
+ variable_names=["x"],
+ parameters={"p": 1},
+ ),
+ operators={
+ 1: [f"identity_{type_name}(x::{type_name}) = x"],
+ 2: [
+ f"""
+ choose_parameter(a::{type_name}, b::{type_name}) = a
+ choose_parameter(a::{type_name}, b::ValidVector) =
+ ValidVector(map(_ -> a, b.x), b.valid)
+ """
+ ],
+ },
+ niterations=10,
+ ncycles_per_iteration=10,
+ populations=2,
+ population_size=20,
+ )
+
+ model.fit(X, y)
+
+ np.testing.assert_array_equal(model.predict(X), y)
+ self.assertEqual(model.get_best()["loss"], 0.0)
+
+ def test_numeric_path_remains_independent(self):
+ X = np.linspace(-1.0, 1.0, 20).reshape(-1, 1)
+ y = X[:, 0]
+ model = PySRRegressor(
+ unary_operators=[],
+ binary_operators=["+"],
+ niterations=1,
+ ncycles_per_iteration=2,
+ populations=1,
+ population_size=8,
+ tournament_selection_n=3,
+ maxsize=7,
+ parallelism="serial",
+ deterministic=True,
+ random_state=0,
+ progress=False,
+ verbosity=0,
+ temp_equation_file=True,
+ should_optimize_constants=False,
+ )
+ model.fit(X, y)
+ prediction = model.predict(X)
+ self.assertEqual(prediction.shape, y.shape)
+ self.assertTrue(np.isfinite(prediction).all())
+
+ model.set_params(
+ type_spec=string_spec(),
+ operators={1: ["identity_value(x::StringValue) = x"]},
+ )
+
+ np.testing.assert_array_equal(model.predict(X), prediction)
+ self.assertTrue(model._supports_export("sympy"))
+
+ def test_configured_full_objective_and_module_local_helpers(self):
+ X, y = string_data()
+ type_name = "ConfiguredObjectiveValue"
+ called = "CONFIGURED_OBJECTIVE_SAW_BATCH"
+ model = tiny_model(
+ string_spec(
+ name=type_name,
+ preamble=(
+ f"const {called} = Ref(false)\n" "const PrivateLoss = Float64"
+ ),
+ loss_type="PrivateLoss",
+ ),
+ operators={1: [f"configured_identity(x::{type_name}) = x"]},
+ elementwise_loss=None,
+ loss_function="""
+ function configured_objective(
+ tree, dataset::Dataset, options, idx=nothing
+ )::PrivateLoss
+ CONFIGURED_OBJECTIVE_SAW_BATCH[] |= idx !== nothing
+ _, complete = eval_tree_array(tree, dataset.X, options)
+ return complete ? 0.0 : Inf
+ end
+ """,
+ batching=True,
+ batch_size=2,
+ complexity_mapping=f"expression -> {type_name} <: Any ? 1 : 1",
+ early_stop_condition="(loss, complexity) -> false",
+ )
+ model.fit(X, y)
+ runtime = model._load_type_spec_runtime()
+
+ self.assertTrue(
+ bool(
+ jl.getindex(
+ jl.getproperty(runtime.configuration_module, jl.Symbol(called))
+ )
+ )
+ )
+ self.assertEqual(model.predict(X).shape, y.shape)
+
+ def test_custom_expression_specs_are_deferred(self):
+ class CustomExpressionSpec(AbstractExpressionSpec):
+ def julia_expression_spec(self):
+ return ExpressionSpec().julia_expression_spec()
+
+ def create_exports(self, *args, **kwargs):
+ return ExpressionSpec().create_exports(*args, **kwargs)
+
+ X, y = string_data()
+ model = tiny_model(
+ string_spec(name="DeferredExpressionValue"),
+ expression_spec=CustomExpressionSpec(),
+ )
+ with self.assertRaisesRegex(
+ ValueError, "CustomExpressionSpec does not support TypeSpec"
+ ):
+ model.fit(X, y)
+
+ def test_rejects_invalid_specs(self):
+ cases = [
+ (dict(name="not an id"), "not an identifier"),
+ (dict(fields={}), "non-empty ordered mapping"),
+ (dict(fields={"not an id": "String"}), "is not an identifier"),
+ (dict(fields={"data": " "}), "requires a Julia type"),
+ (dict(sample=" "), "must contain Julia source"),
+ (dict(scalar_constants="value -> [1.0]"), "must be provided together"),
+ (dict(mutate=None), "requires an explicit `mutate`"),
+ (dict(string=" "), "cannot be empty"),
+ ]
+ for overrides, message in cases:
+ with self.subTest(**overrides):
+ with self.assertRaisesRegex(ValueError, message):
+ string_spec(**overrides)
+
+ def test_rejects_invalid_configurations(self):
+ no_loss = dict(
+ elementwise_loss=None, loss_function=None, loss_function_expression=None
+ )
+ operators = {1: ["x -> x"]}
+ cases = [
+ (string_spec(), operators, no_loss, "exactly one of"),
+ (
+ string_spec(loss_type="Float64"),
+ operators,
+ {**no_loss, "elementwise_loss": "(p, t) -> 0.0"},
+ "return type is inferred",
+ ),
+ (
+ string_spec(),
+ operators,
+ {**no_loss, "loss_function": "f(tree, dataset, options) = 0.0"},
+ "explicit `loss_type`",
+ ),
+ (string_spec(), None, no_loss, "requires explicit"),
+ (string_spec(), {0: ["x -> x"]}, no_loss, "positive integers"),
+ (string_spec(), {1: []}, no_loss, "cannot be empty"),
+ (string_spec(), {1: [" "]}, no_loss, "must contain Julia source"),
+ ]
+ for spec, operator_table, losses, message in cases:
+ with self.subTest(message=message):
+ with self.assertRaisesRegex(ValueError, message):
+ validate_type_spec_configuration(spec, operator_table, **losses)
+
+ def test_fit_data_validation(self):
+ X = np.array([["a"], ["b"]], dtype=object)
+ y = X[:, 0].copy()
+
+ def prepare(
+ model=None,
+ X=X,
+ y=y,
+ Xresampled=None,
+ weights=None,
+ variable_names=None,
+ X_units=None,
+ y_units=None,
+ ):
+ return prepare_type_spec_fit_data(
+ model or PySRRegressor(),
+ X,
+ y,
+ Xresampled,
+ weights,
+ variable_names,
+ None,
+ X_units,
+ y_units,
+ )
+
+ with self.assertRaisesRegex(NotImplementedError, "denoising"):
+ prepare(model=PySRRegressor(denoise=True))
+ with self.assertRaisesRegex(ValueError, "2D array"):
+ prepare(X=np.empty((1, 1, 1), dtype=object))
+ with self.assertRaisesRegex(NotImplementedError, "one output"):
+ prepare(y=np.empty((2, 2), dtype=object))
+ with self.assertRaisesRegex(ValueError, "inconsistent numbers of samples"):
+ prepare(y=y[:1])
+ with self.assertRaisesRegex(ValueError, "at least one sample"):
+ prepare(X=np.empty((0, 1), dtype=object), y=np.empty(0, dtype=object))
+ with self.assertRaisesRegex(NotImplementedError, "weights"):
+ prepare(weights=np.ones(2))
+ with self.assertRaisesRegex(NotImplementedError, "units"):
+ prepare(X_units=["m"])
+ with self.assertWarnsRegex(UserWarning, "reset to `None`"):
+ prepare(X=pd.DataFrame({"a": ["a", "b"]}), variable_names=["a"])
+ model = PySRRegressor()
+ with self.assertWarnsRegex(UserWarning, "Spaces"):
+ prepare(model=model, y=y.reshape(-1, 1), variable_names=["a b"])
+ self.assertEqual(list(model.feature_names_in_), ["a_b"])
+
+ def test_prediction_data_validation(self):
+ model = PySRRegressor()
+ model.n_features_in_ = 1
+ model.selection_mask_ = None
+ model.feature_names_in_ = np.array(["a_b"])
+ with self.assertRaisesRegex(ValueError, "2D array"):
+ prepare_type_spec_prediction_data(model, np.empty((1, 1, 1), dtype=object))
+ with self.assertRaisesRegex(ValueError, "different number of features"):
+ prepare_type_spec_prediction_data(
+ model, np.array([["a", "b"]], dtype=object)
+ )
+ with self.assertRaisesRegex(ValueError, "missing features"):
+ prepare_type_spec_prediction_data(model, pd.DataFrame({"z": ["a"]}))
+ with self.assertWarnsRegex(UserWarning, "Spaces"):
+ out = prepare_type_spec_prediction_data(model, pd.DataFrame({"a b": ["a"]}))
+ self.assertEqual(out.shape, (1, 1))
+
+ model.n_features_in_ = 2
+ model.selection_mask_ = np.array([True, False])
+ out = prepare_type_spec_prediction_data(
+ model, np.array([["a", "b"]], dtype=object)
+ )
+ np.testing.assert_array_equal(out, np.array([["a"]], dtype=object))
+
+ def test_julia_array_conversion_errors(self):
+ runtime = load_type_spec_runtime(
+ compile_type_spec(string_spec(name="ConversionErrorValue"))
+ )
+ with self.assertRaisesRegex(ValueError, "1D or 2D"):
+ type_spec_to_julia_array(runtime, np.empty((1, 1, 1), dtype=object))
+ with self.assertRaises(ValueError):
+ type_spec_to_julia_array(runtime, np.array([1.0, 2.0]))
+
+ def test_runtime_load_rejects_inconsistent_definition_values(self):
+ definition = compile_type_spec_runtime(
+ string_spec(name="InconsistentLoadValue"),
+ {1: ["x -> x"]},
+ elementwise_loss="(prediction, target) -> 0.0",
+ loss_function=None,
+ loss_function_expression=None,
+ complexity_mapping=None,
+ early_stop_condition=None,
+ )
+ load_type_spec_runtime(definition)
+ with patch("pysr.type_specs._runtime_sources", return_value=["x -> x"] * 99):
+ with self.assertRaisesRegex(RuntimeError, "inconsistent"):
+ load_type_spec_runtime(definition)
+
+ def test_fitted_model_guards(self):
+ X, y = string_data()
+ spec = string_spec()
+ model = tiny_model(spec)
+ model.fit(X, y)
+
+ with self.assertRaises(NotImplementedError):
+ model.score(X, y)
+ with self.assertRaises(ValueError):
+ model.predict(np.array([[1.0]]))
+
+ model.set_params(warm_start=True, type_spec=None)
+ with self.assertRaisesRegex(ValueError, "Cannot warm-start"):
+ model.fit(X, y)
+
+ model.set_params(warm_start=False, type_spec=spec)
+ with warnings.catch_warnings():
+ warnings.simplefilter("ignore")
+ with self.assertRaisesRegex(NotImplementedError, "weights"):
+ model.fit(X, y, weights=np.ones(len(y)))
+ self.assertFalse(hasattr(model, "_type_spec_runtime_definition_"))
+
+ def test_guess_seeds_the_search_with_a_custom_constant(self):
+ type_name = "GuessVectorValue"
+ model = tiny_model(
+ vector_spec(name=type_name),
+ operators={
+ 2: [
+ f"add_vectors(a::{type_name}, b::{type_name}) = "
+ f"{type_name}(a.data + b.data)"
+ ]
+ },
+ elementwise_loss=(
+ f"vector_loss(a::{type_name}, b::{type_name})::Float64 = "
+ "sum(abs2, a.data - b.data)"
+ ),
+ guesses=[f"add_vectors(x0, {type_name}([1.5, -2.5]))"],
+ )
+ X = np.empty((4, 1), dtype=object)
+ X[:, 0] = [np.array([1.0, 2.0])] * 4
+ y = np.empty(4, dtype=object)
+ y[:] = [np.array([2.5, -0.5])] * 4
+
+ model.fit(X, y)
+
+ self.assertEqual(model.equations_.loss.min(), 0.0)
+ self.assertIn("[1.5, -2.5]", " ".join(model.equations_.equation))
+
+ def test_template_guess_seeds_each_expression(self):
+ type_name = "GuessTemplateValue"
+ model = tiny_model(
+ string_spec(name=type_name),
+ expression_spec=TemplateExpressionSpec(
+ combine="f(x)",
+ expressions=["f"],
+ variable_names=["x"],
+ ),
+ operators={
+ 2: [
+ f"concat(a::{type_name}, b::{type_name}) = "
+ f"{type_name}(a.data * b.data)"
+ ]
+ },
+ elementwise_loss=(
+ f"string_loss(a::{type_name}, b::{type_name})::Float64 = "
+ "a.data == b.data ? 0.0 : 1.0"
+ ),
+ guesses=[{"f": f'concat(#1, {type_name}("!"))'}],
+ )
+ X = np.array([["a"], ["b"], ["a"], ["b"]], dtype=object)
+ y = np.array(["a!", "b!", "a!", "b!"], dtype=object)
+
+ model.fit(X, y, variable_names=["x"])
+
+ self.assertEqual(model.equations_.loss.min(), 0.0)
+ self.assertIn('concat(#1, "!")', " ".join(model.equations_.equation))
+
+ def test_template_parameter_guess_with_string_values(self):
+ type_name = "GuessTemplateParameterString"
+ model = tiny_model(
+ string_spec(
+ name=type_name,
+ sample=f'rng -> {type_name}("a")',
+ mutate="(rng, value, temperature) -> value",
+ ),
+ expression_spec=TemplateExpressionSpec(
+ combine="choose_parameter(p[1], f(x))",
+ expressions=["f"],
+ variable_names=["x"],
+ parameters={"p": 1},
+ ),
+ operators={
+ 1: [f"identity_{type_name}(x::{type_name}) = x"],
+ 2: [
+ f"""
+ choose_parameter(a::{type_name}, b::{type_name}) = a
+ choose_parameter(a::{type_name}, b::ValidVector) =
+ ValidVector(map(_ -> a, b.x), b.valid)
+ """
+ ],
+ },
+ guesses=[{"f": "#1", "p": ["!"]}],
+ )
+ X = np.array([["a"], ["b"], ["a"], ["b"]], dtype=object)
+ y = np.full(len(X), "!", dtype=object)
+
+ model.fit(X, y, variable_names=["x"])
+
+ np.testing.assert_array_equal(model.predict(X), y)
+
+ def test_template_parameter_guess_with_vector_values(self):
+ type_name = "GuessTemplateParameterVector"
+ parameter_values = object_array_1d([np.array([1.5, -2.5])])
+ original_value = parameter_values[0].copy()
+ model = tiny_model(
+ vector_spec(name=type_name),
+ expression_spec=TemplateExpressionSpec(
+ combine="choose_parameter(p[1], f(x))",
+ expressions=["f"],
+ variable_names=["x"],
+ parameters={"p": 1},
+ ),
+ operators={
+ 1: [f"identity_{type_name}(x::{type_name}) = x"],
+ 2: [
+ f"""
+ choose_parameter(a::{type_name}, b::{type_name}) = a
+ choose_parameter(a::{type_name}, b::ValidVector) =
+ ValidVector(map(_ -> a, b.x), b.valid)
+ """
+ ],
+ },
+ guesses=[{"f": "#1", "p": parameter_values}],
+ )
+ X = np.empty((4, 1), dtype=object)
+ X[:, 0] = [np.array([1.0, 2.0])] * 4
+ y = np.empty(4, dtype=object)
+ y[:] = [np.array([1.5, -2.5])] * 4
+
+ model.fit(X, y, variable_names=["x"])
+
+ self.assertEqual(
+ [list(value) for value in model.predict(X)],
+ [list(value) for value in y],
+ )
+ np.testing.assert_array_equal(parameter_values[0], original_value)
+
+ def test_template_parameter_guess_with_packet_values(self):
+ type_name = "GuessTemplateParameterPacket"
+ model = tiny_model(
+ TypeSpec(
+ type_name,
+ fields={"number": "Float64", "label": "String"},
+ sample=f'rng -> {type_name}(0.0, "")',
+ mutate="(rng, value, temperature) -> value",
+ ),
+ expression_spec=TemplateExpressionSpec(
+ combine="choose_parameter(p[1], f(x))",
+ expressions=["f"],
+ variable_names=["x"],
+ parameters={"p": 1},
+ ),
+ operators={
+ 1: [f"identity_{type_name}(x::{type_name}) = x"],
+ 2: [
+ f"""
+ choose_parameter(a::{type_name}, b::{type_name}) = a
+ choose_parameter(a::{type_name}, b::ValidVector) =
+ ValidVector(map(_ -> a, b.x), b.valid)
+ """
+ ],
+ },
+ guesses=[{"f": "#1", "p": [(2.0, "two")]}],
+ )
+ X = np.empty((4, 1), dtype=object)
+ X[:, 0] = object_array_1d([(1.0, "one")] * 4)
+ y = object_array_1d([(2.0, "two")] * 4)
+
+ model.fit(X, y, variable_names=["x"])
+
+ self.assertEqual(model.predict(X).tolist(), y.tolist())
+
+ def test_guess_rejects_a_constant_that_is_not_the_custom_type(self):
+ type_name = "GuessConversionValue"
+ model = tiny_model(
+ vector_spec(name=type_name),
+ guesses=[f"identity_{type_name}(2.0)"],
+ )
+ X = np.empty((4, 1), dtype=object)
+ X[:, 0] = [np.array([1.0, 2.0])] * 4
+ y = np.empty(4, dtype=object)
+ y[:] = [np.array([1.0, 2.0])] * 4
+
+ with self.assertRaisesRegex(
+ JuliaError, f"Cannot `convert` an object of type Float64 to .*{type_name}"
+ ):
+ model.fit(X, y)
+
+ def test_guess_rejects_an_undefined_constant_constructor(self):
+ model = tiny_model(
+ string_spec(),
+ guesses=['MissingGuessValue("!")'],
+ )
+ X, y = string_data()
+
+ with self.assertRaisesRegex(
+ JuliaError, "Tried to interpolate function `MissingGuessValue`"
+ ):
+ model.fit(X, y)
+
+
+if __name__ == "__main__":
+ unittest.main()
diff --git a/pysr/type_specs.py b/pysr/type_specs.py
new file mode 100644
index 000000000..351b16229
--- /dev/null
+++ b/pysr/type_specs.py
@@ -0,0 +1,1368 @@
+from __future__ import annotations
+
+import copy
+import hashlib
+import json
+import warnings
+from dataclasses import dataclass
+from functools import cache
+from textwrap import dedent
+from typing import Any
+
+import numpy as np
+import pandas as pd
+from juliacall import JuliaError # type: ignore
+
+from .julia_helpers import jl_array
+from .julia_import import AnyValue, SymbolicRegression, jl
+
+_TYPE_MODULE_INSTALLED = "_PYSR_TYPE_SPEC_INSTALLED"
+
+
+def object_array_1d(values: Any) -> np.ndarray:
+ """Build a 1D object array whose cells are the logical values themselves."""
+ if isinstance(values, np.ndarray):
+ return values if values.dtype == object else values.astype(object)
+ values = list(values)
+ array = np.empty(len(values), dtype=object)
+ for i, value in enumerate(values):
+ array[i] = value
+ return array
+
+
+def object_array_2d(values: Any) -> np.ndarray:
+ """Build a 2D object array whose cells are the logical values themselves."""
+ if isinstance(values, np.ndarray):
+ return values if values.dtype == object else values.astype(object)
+ try:
+ values = list(values)
+ if any(isinstance(row, (str, bytes)) for row in values):
+ raise TypeError
+ rows = [list(row) for row in values]
+ except TypeError:
+ raise ValueError("TypeSpec X must be a 2D array of logical values.")
+ n_columns = len(rows[0]) if rows else 0
+ array = np.empty((len(rows), n_columns), dtype=object)
+ for i, row in enumerate(rows):
+ if len(row) != n_columns:
+ raise ValueError("All rows of X must have the same number of features.")
+ for j, value in enumerate(row):
+ array[i, j] = value
+ return array
+
+
+@dataclass
+class TypeSpec:
+ """Definition of a custom value type for symbolic regression.
+
+ Parameters
+ ----------
+ name : str
+ Name of the generated Julia type. Use this name in operator and hook
+ definitions.
+ fields : dict[str, str]
+ Ordered mapping from field names to Julia field types.
+ sample : str
+ Julia callable with signature ``rng -> (...)::{name}``.
+ scalar_constants : str, optional
+ Julia callable with signature ``value::{name} -> (...)::Vector{Float64}``. Provide
+ together with ``with_scalar_constants`` to enable continuous constant
+ optimization.
+ with_scalar_constants : str, optional
+ Julia callable with signature ``(value::{name}, scalar_constants::Vector{Float64}) -> (...)::{name}``.
+ init : str, optional
+ Julia callable with signature ``() -> value::{name}``. The default repurposes `sample`
+ using a deterministic local random-number generator.
+ mutate : str, optional
+ Julia callable with signature ``(rng, value::{name}, temperature::Float64) -> value::{name}``.
+ Use this to mutate a value of the type.
+ The `temperature` (from simulated annealing) is between 0 and 1. High temperatures
+ generally should be used for more aggressive mutations, and low temperatures for
+ more conservative mutations.
+ The default mutation repurposes `scalar_constants` and `with_scalar_constants`
+ to mutate the set with SymbolicRegression.jl's default constant mutation,
+ and put it back into a new value.
+ is_valid : str, optional
+ Julia callable with signature ``value::{name} -> (...)::Bool``.
+ Use this to check if your type is in a valid state. For example, if your type
+ has been processed in a way that invalidates it (like a NaN, but specific to your type),
+ you may use this function to check for that. This will be used to quit evaluation early.
+ The default checks that every scalar constant is finite, or accepts every value for a
+ non-optimizable type.
+ init_invalid : str, optional
+ Julia callable with signature ``() -> value::{name}`` that constructs an
+ invalid value. The result must have this type and fail ``is_valid``.
+ Used by DynamicExpressions to fill failed convenience evaluations.
+ Omit this hook for types with no invalid state. It receives no prototype
+ or shape information.
+ string : str, optional
+ Julia callable with signature ``value::{name} -> (...)::AbstractString`` used to print
+ values.
+ preamble : str, optional
+ Julia source evaluated once before the generated type definition. Types
+ and functions it defines are visible to the hooks and to operator and
+ objective sources; imports are not, so any source needing a package
+ beyond SymbolicRegression must import it itself.
+ definitions : str, optional
+ Julia source evaluated once directly after the generated type, for
+ constructors and methods that mention it.
+ loss_type : str, optional
+ Concrete Julia ``AbstractFloat`` type returned by a custom full
+ objective. Elementwise loss return types are inferred.
+ """
+
+ name: str
+ fields: dict[str, str]
+ sample: str
+ scalar_constants: str | None = None
+ with_scalar_constants: str | None = None
+ init: str | None = None
+ mutate: str | None = None
+ is_valid: str | None = None
+ string: str | None = None
+ preamble: str | None = None
+ definitions: str | None = None
+ loss_type: str | None = None
+ init_invalid: str | None = None
+
+ def __post_init__(self) -> None:
+ if not isinstance(self.name, str) or not self.name.isidentifier():
+ raise ValueError(f"TypeSpec name {self.name!r} is not an identifier.")
+ if not isinstance(self.fields, dict) or not self.fields:
+ raise ValueError("`fields` must be a non-empty ordered mapping.")
+ for name, field_type in self.fields.items():
+ if not isinstance(name, str) or not name.isidentifier():
+ raise ValueError(f"TypeSpec field name {name!r} is not an identifier.")
+ if not isinstance(field_type, str) or not field_type.strip():
+ raise ValueError(f"TypeSpec field `{name}` requires a Julia type.")
+ if not isinstance(self.sample, str) or not self.sample.strip():
+ raise ValueError("`sample` must contain Julia source.")
+
+ if (self.scalar_constants is None) != (self.with_scalar_constants is None):
+ raise ValueError(
+ "`scalar_constants` and `with_scalar_constants` must be provided "
+ "together."
+ )
+ if not self.can_optimize and self.mutate is None:
+ raise ValueError(
+ "A non-optimizable TypeSpec requires an explicit `mutate` callable."
+ )
+
+ for name in (
+ "scalar_constants",
+ "with_scalar_constants",
+ "init",
+ "mutate",
+ "is_valid",
+ "init_invalid",
+ "string",
+ "preamble",
+ "definitions",
+ "loss_type",
+ ):
+ value = getattr(self, name)
+ if value is not None and (not isinstance(value, str) or not value.strip()):
+ raise ValueError(f"`{name}` cannot be empty.")
+
+ @property
+ def can_optimize(self) -> bool:
+ return self.scalar_constants is not None
+
+
+@dataclass(frozen=True)
+class _TypeSpecDefinition:
+ spec: TypeSpec
+ fingerprint: str
+ source: str
+
+ @property
+ def module_name(self) -> str:
+ return f"_PySRTypeSpec_{self.fingerprint[:20]}"
+
+
+@dataclass(frozen=True)
+class _TypeSpecRuntimeDefinition:
+ type_spec: _TypeSpecDefinition
+ operators: tuple[tuple[int, tuple[str, ...]], ...]
+ elementwise_loss: str | None
+ loss_function: str | None
+ loss_function_expression: str | None
+ complexity_mapping: str | None
+ early_stop_condition: str | None
+ expression_spec: str
+ expression_spec_function_selector: str | None
+ fingerprint: str
+ source: str
+
+ @property
+ def module_name(self) -> str:
+ return f"_PySRConfig_{self.fingerprint[:20]}"
+
+
+@dataclass(frozen=True)
+class _TypeSpecRuntime:
+ definition: _TypeSpecDefinition | _TypeSpecRuntimeDefinition
+ module: AnyValue
+ value_type: AnyValue
+ configuration_module: AnyValue
+ operator_functions: dict[int, tuple[AnyValue, ...]]
+ operator_names: dict[int, tuple[str, ...]]
+ elementwise_loss: AnyValue | None = None
+ loss_function: AnyValue | None = None
+ loss_function_expression: AnyValue | None = None
+ complexity_mapping: AnyValue | None = None
+ early_stop_condition: float | AnyValue | None = None
+ expression_spec: AnyValue | None = None
+
+ @property
+ def type_definition(self) -> _TypeSpecDefinition:
+ if isinstance(self.definition, _TypeSpecRuntimeDefinition):
+ return self.definition.type_spec
+ return self.definition
+
+ @property
+ def spec(self) -> TypeSpec:
+ return self.type_definition.spec
+
+
+def _quoted(source: str) -> str:
+ return json.dumps(source, ensure_ascii=False).replace("$", r"\$")
+
+
+def _block(source: str) -> str:
+ return dedent(source).strip()
+
+
+def _runtime_sources(
+ operators: tuple[tuple[int, tuple[str, ...]], ...],
+ elementwise_loss: str | None,
+ loss_function: str | None,
+ loss_function_expression: str | None,
+ complexity_mapping: str | None,
+ early_stop_condition: str | None,
+ expression_spec: str,
+) -> tuple[str, ...]:
+ sources = tuple(
+ source for _, operator_sources in operators for source in operator_sources
+ )
+ optional_sources = (
+ elementwise_loss,
+ loss_function,
+ loss_function_expression,
+ complexity_mapping,
+ early_stop_condition,
+ )
+ return (
+ *sources,
+ *(source for source in optional_sources if source is not None),
+ expression_spec,
+ )
+
+
+def _runtime_install_source(
+ type_definition: _TypeSpecDefinition,
+ runtime_module_name: str,
+ sources: tuple[str, ...],
+ expression_spec_function_selector: str | None,
+) -> str:
+ source_literals = ",\n".join(f" {_quoted(source)}" for source in sources)
+ selector = _optional_source(expression_spec_function_selector)
+ return _block(f"""
+ let type_module_name = Symbol({_quoted(type_definition.module_name)})
+ installed =
+ isdefined(Main, type_module_name) &&
+ isdefined(
+ Base.invokelatest(getproperty, Main, type_module_name),
+ Symbol({_quoted(_TYPE_MODULE_INSTALLED)}),
+ )
+ if !installed
+ Base.include_string(
+ Main,
+ {_quoted(type_definition.source)},
+ {_quoted("PySR." + type_definition.module_name)},
+ )
+ end
+ end
+ let
+ parent = getproperty(
+ Main, Symbol({_quoted(type_definition.module_name)})
+ )
+ module_name = Symbol({_quoted(runtime_module_name)})
+ needs_install = !isdefined(parent, module_name)
+ if !needs_install
+ existing = Base.invokelatest(getproperty, parent, module_name)
+ needs_install = !isdefined(existing, :_definition_values)
+ end
+ if needs_install
+ Core.eval(parent, Expr(:module, true, module_name, Expr(:block)))
+ module_ = Base.invokelatest(getproperty, parent, module_name)
+ Core.eval(module_, :(using SymbolicRegression))
+ for name in getproperty(parent, :_PYSR_PARENT_BINDING_NAMES)
+ isdefined(module_, name) && continue
+ Core.eval(
+ module_,
+ Expr(
+ :const,
+ Expr(
+ :(=),
+ name,
+ QuoteNode(getproperty(parent, name)),
+ ),
+ ),
+ )
+ end
+ sources = String[
+ {source_literals}
+ ]
+ values = Any[]
+ for (index, source) in enumerate(sources)
+ push!(
+ values,
+ Base.include_string(
+ module_,
+ source,
+ "PySR TypeSpec configuration $index",
+ ),
+ )
+ end
+ selector_source = {selector}
+ if selector_source !== nothing
+ selector_value = Base.include_string(
+ module_,
+ selector_source,
+ "PySR TypeSpec expression function selector",
+ )
+ push!(values, Base.invokelatest(selector_value, values[end]))
+ end
+ Core.eval(
+ module_,
+ Expr(
+ :const,
+ Expr(:(=), :_definition_values, QuoteNode(Tuple(values))),
+ ),
+ )
+ end
+ end
+ """) + "\n"
+
+
+def compile_type_spec_runtime(
+ spec: TypeSpec,
+ operators: dict[int, list[str]],
+ *,
+ elementwise_loss: str | None,
+ loss_function: str | None,
+ loss_function_expression: str | None,
+ complexity_mapping: str | None,
+ early_stop_condition: str | None,
+ expression_spec: str | None = None,
+ expression_spec_function_selector: str | None = None,
+) -> _TypeSpecRuntimeDefinition:
+ validate_type_spec_configuration(
+ spec,
+ operators,
+ elementwise_loss=elementwise_loss,
+ loss_function=loss_function,
+ loss_function_expression=loss_function_expression,
+ )
+ type_definition = compile_type_spec(spec)
+ normalized_operators = tuple(
+ (arity, tuple(sources)) for arity, sources in _normalize_operators(operators)
+ )
+ if expression_spec is None:
+ expression_spec = _block(f"""
+ SymbolicRegression.ExpressionSpec(
+ node_type=SymbolicRegression.InterfaceDynamicExpressionsModule.DE.Node{{
+ {spec.name}
+ }},
+ )
+ """)
+ payload = {
+ "type_spec": type_definition.fingerprint,
+ "operators": normalized_operators,
+ "elementwise_loss": elementwise_loss,
+ "loss_function": loss_function,
+ "loss_function_expression": loss_function_expression,
+ "complexity_mapping": complexity_mapping,
+ "early_stop_condition": early_stop_condition,
+ "expression_spec": expression_spec,
+ "expression_spec_function_selector": expression_spec_function_selector,
+ "loss_type": spec.loss_type,
+ }
+ fingerprint = hashlib.sha256(
+ json.dumps(payload, sort_keys=True).encode()
+ ).hexdigest()
+ runtime_module_name = f"_PySRConfig_{fingerprint[:20]}"
+ sources = _runtime_sources(
+ normalized_operators,
+ elementwise_loss,
+ loss_function,
+ loss_function_expression,
+ complexity_mapping,
+ early_stop_condition,
+ expression_spec,
+ )
+ source = _runtime_install_source(
+ type_definition,
+ runtime_module_name,
+ sources,
+ expression_spec_function_selector,
+ )
+ return _TypeSpecRuntimeDefinition(
+ type_spec=type_definition,
+ operators=normalized_operators,
+ elementwise_loss=elementwise_loss,
+ loss_function=loss_function,
+ loss_function_expression=loss_function_expression,
+ complexity_mapping=complexity_mapping,
+ early_stop_condition=early_stop_condition,
+ expression_spec=expression_spec,
+ expression_spec_function_selector=expression_spec_function_selector,
+ fingerprint=fingerprint,
+ source=source,
+ )
+
+
+def compile_type_spec_runtime_for_model(
+ model: Any,
+ operators: dict[int, list[str]],
+) -> _TypeSpecRuntimeDefinition:
+ """Compile one model configuration without evaluating its Julia sources."""
+ from .expression_specs import ExpressionSpec
+
+ spec = model.type_spec
+ assert spec is not None
+ expression_spec = model.expression_spec_
+ expression_spec_source = expression_spec._julia_expression_spec_source(
+ prototype=f"SymbolicRegression.init_value({spec.name})"
+ )
+ if expression_spec_source is None and type(expression_spec) is not ExpressionSpec:
+ raise ValueError(
+ f"{type(expression_spec).__name__} declares TypeSpec support "
+ "but does not provide Julia source."
+ )
+ return compile_type_spec_runtime(
+ spec,
+ operators,
+ elementwise_loss=model.elementwise_loss,
+ loss_function=model.loss_function,
+ loss_function_expression=model.loss_function_expression,
+ complexity_mapping=model.complexity_mapping,
+ early_stop_condition=(
+ None
+ if model.early_stop_condition is None
+ else str(model.early_stop_condition)
+ ),
+ expression_spec=expression_spec_source,
+ expression_spec_function_selector=(
+ expression_spec._julia_expression_spec_function_selector()
+ ),
+ )
+
+
+_TYPE_SPEC_MODULE = _block(r"""
+ import SymbolicRegression: init_value, mutate_value, sample_value
+ import SymbolicRegression.ConstantOptimizationModule: can_optimize
+ import SymbolicRegression.InterfaceDynamicExpressionsModule: string_constant
+ import SymbolicRegression.InterfaceDynamicExpressionsModule.DE:
+ count_scalar_constants, get_number_type, is_valid, invalid_value,
+ pack_scalar_constants!, unpack_scalar_constants
+ import SymbolicRegression.InterfaceDynamicExpressionsModule.DE.StringsModule:
+ needs_brackets
+
+ const _config = __TYPE_SPEC_CONFIG__
+ abstract type _TypeSpecValue end
+
+ macro _define_type_spec(config_expression)
+ config = Core.eval(__module__, config_expression)
+ fields = map(config.fields) do (name, type)
+ :($(Symbol(name))::$(Meta.parse(type)))
+ end
+ type_definition = Expr(
+ :struct,
+ false,
+ Expr(:<:, Symbol(config.name), :_TypeSpecValue),
+ Expr(:block, fields...),
+ )
+ return esc(type_definition)
+ end
+
+ _include(source, label) = Base.include_string(@__MODULE__, source, label)
+ _config.preamble === nothing ||
+ _include(_config.preamble, "TypeSpec.preamble")
+ @_define_type_spec _config
+ const _value_type = getfield(@__MODULE__, Symbol(_config.name))
+ _config.definitions === nothing ||
+ _include(_config.definitions, "TypeSpec.definitions")
+ SymbolicRegression.parse_scope(::Type{<:_TypeSpecValue}) = @__MODULE__
+
+ _fields(value) = ntuple(i -> getfield(value, i), fieldcount(_value_type))
+ Base.:(==)(a::_TypeSpecValue, b::_TypeSpecValue) = _fields(a) == _fields(b)
+ Base.isequal(a::_TypeSpecValue, b::_TypeSpecValue) =
+ isequal(_fields(a), _fields(b))
+ Base.hash(value::_TypeSpecValue, h::UInt) = hash(_fields(value), h)
+
+ const _sample = _include(_config.sample, "TypeSpec.sample")
+ const _init = _config.init === nothing ?
+ () -> _sample(Random.Xoshiro(0)) :
+ _include(_config.init, "TypeSpec.init")
+ init_value(::Type{<:_TypeSpecValue}) = _init()
+ sample_value(rng::AbstractRNG, ::Type{<:_TypeSpecValue}, options) = _sample(rng)
+ can_optimize(::Type{<:_TypeSpecValue}, _) = _config.optimizable
+
+ if _config.optimizable
+ const _scalar_constants =
+ _include(_config.scalar_constants, "TypeSpec.scalar_constants")
+ const _with_scalar_constants = _include(
+ _config.with_scalar_constants,
+ "TypeSpec.with_scalar_constants",
+ )
+ get_number_type(::Type{<:_TypeSpecValue}) =
+ eltype(_scalar_constants(init_value(_value_type)))
+ end
+
+ count_scalar_constants(value::_TypeSpecValue) =
+ _config.optimizable ? length(_scalar_constants(value)) : 0
+
+ if _config.optimizable
+ function pack_scalar_constants!(
+ buffer::AbstractVector{<:Number}, idx::Int, value::_TypeSpecValue
+ )
+ scalar_constants = _scalar_constants(value)
+ copyto!(
+ buffer,
+ idx,
+ scalar_constants,
+ firstindex(scalar_constants),
+ length(scalar_constants),
+ )
+ return idx + length(scalar_constants)
+ end
+ function unpack_scalar_constants(
+ buffer::AbstractVector{<:Number}, idx::Int, value::_TypeSpecValue
+ )
+ count = length(_scalar_constants(value))
+ scalar_constants = @view buffer[idx:(idx + count - 1)]
+ return idx + count, _with_scalar_constants(value, scalar_constants)
+ end
+ end
+
+ const _mutate = if _config.mutate === nothing
+ function (rng, value, temperature, mutation)
+ scalar_constants = collect(_scalar_constants(value))
+ isempty(scalar_constants) && return _sample(rng)
+ i = rand(rng, eachindex(scalar_constants))
+ scalar_constants[i] =
+ SymbolicRegression.MutationFunctionsModule.mutate_value(
+ rng, scalar_constants[i], temperature, mutation
+ )
+ return _with_scalar_constants(value, scalar_constants)
+ end
+ else
+ mutate = _include(_config.mutate, "TypeSpec.mutate")
+ (rng, value, temperature, _) -> mutate(rng, value, temperature)
+ end
+ mutate_value(
+ rng::AbstractRNG,
+ value::_TypeSpecValue,
+ temperature,
+ mutation::SymbolicRegression.ConstantMutation,
+ ) = _mutate(rng, value, temperature, mutation)
+
+ const _is_valid = if _config.is_valid !== nothing
+ _include(_config.is_valid, "TypeSpec.is_valid")
+ elseif _config.optimizable
+ value -> all(isfinite, _scalar_constants(value))
+ else
+ _ -> true
+ end
+ is_valid(value::_TypeSpecValue) = _is_valid(value)
+
+ if _config.init_invalid !== nothing
+ const _init_invalid = _include(_config.init_invalid, "TypeSpec.init_invalid")
+ invalid_value(::Type{<:_TypeSpecValue}) = _init_invalid()
+ end
+
+ const _string = if _config.string === nothing
+ function (value)
+ fields = map(1:fieldcount(_value_type)) do i
+ sprint(show, getfield(value, i); context=:compact => true)
+ end
+ return fieldcount(_value_type) == 1 ? only(fields) :
+ string(_config.name, "(", join(fields, ", "), ")")
+ end
+ else
+ _include(_config.string, "TypeSpec.string")
+ end
+ Base.show(io::IO, value::_TypeSpecValue) = print(io, _string(value))
+ needs_brackets(::_TypeSpecValue) = false
+ string_constant(
+ value::_TypeSpecValue, ::Val{precision}, unit
+ ) where {precision} = _string(value) * unit
+
+ function _convert_value(x)
+ x isa _value_type && return x
+ x = PythonCall.Py(x)
+ fieldcount(_value_type) == 1 ||
+ PythonCall.pyhasattr(x, "__len__") &&
+ PythonCall.pylen(x) == fieldcount(_value_type) ||
+ throw(ArgumentError(
+ "TypeSpec values require exactly $(fieldcount(_value_type)) fields."
+ ))
+ values = ntuple(fieldcount(_value_type)) do i
+ source = fieldcount(_value_type) == 1 ? x : x[i - 1]
+ PythonCall.pyconvert(fieldtype(_value_type, i), source)
+ end
+ return _value_type(values...)
+ end
+ _convert_array(values, dims) = reshape(
+ _value_type[_convert_value(value) for value in values], Tuple(dims)
+ )
+ """)
+
+
+def _optional_source(source: str | None) -> str:
+ return "nothing" if source is None else _quoted(source)
+
+
+def validate_type_spec_loss_configuration(
+ spec: TypeSpec,
+ elementwise_loss: str | None,
+ loss_function: str | None,
+ loss_function_expression: str | None,
+) -> str:
+ configured = [
+ ("elementwise_loss", elementwise_loss),
+ ("loss_function", loss_function),
+ ("loss_function_expression", loss_function_expression),
+ ]
+ selected = [mode for mode, source in configured if source is not None]
+ if len(selected) != 1:
+ raise ValueError(
+ "TypeSpec requires exactly one of `elementwise_loss`, `loss_function`, "
+ "and `loss_function_expression`."
+ )
+ mode = selected[0]
+ if mode == "elementwise_loss" and spec.loss_type is not None:
+ raise ValueError(
+ "Do not set `loss_type` with `elementwise_loss`; its return type is inferred."
+ )
+ if mode != "elementwise_loss" and spec.loss_type is None:
+ raise ValueError("TypeSpec full objectives require an explicit `loss_type`.")
+ return mode
+
+
+def _normalize_operators(
+ operators: dict[int, list[str]] | None,
+) -> list[tuple[int, list[str]]]:
+ if not operators:
+ raise ValueError("TypeSpec requires explicit `operators={...}`.")
+ normalized = []
+ for arity, sources in operators.items():
+ if type(arity) is not int or arity < 1:
+ raise ValueError("TypeSpec operator arities must be positive integers.")
+ if not sources:
+ raise ValueError(f"TypeSpec operator arity {arity} cannot be empty.")
+ if any(not isinstance(source, str) or not source.strip() for source in sources):
+ raise ValueError("Every TypeSpec operator must contain Julia source.")
+ normalized.append((arity, list(sources)))
+ return sorted(normalized)
+
+
+def validate_type_spec_configuration(
+ spec: TypeSpec,
+ operators: dict[int, list[str]] | None,
+ *,
+ elementwise_loss: str | None,
+ loss_function: str | None,
+ loss_function_expression: str | None,
+) -> None:
+ _normalize_operators(operators)
+ validate_type_spec_loss_configuration(
+ spec, elementwise_loss, loss_function, loss_function_expression
+ )
+
+
+def compile_type_spec(spec: TypeSpec) -> _TypeSpecDefinition:
+ """Create deterministic Julia source without evaluating user code."""
+ field_sources = ",\n".join(
+ f" {_quoted(name)} => {_quoted(field_type)}"
+ for name, field_type in spec.fields.items()
+ )
+ config = _block(f"""
+ (
+ name = {_quoted(spec.name)},
+ fields = (
+ {field_sources},
+ ),
+ sample = {_quoted(spec.sample)},
+ scalar_constants = {_optional_source(spec.scalar_constants)},
+ with_scalar_constants = {_optional_source(spec.with_scalar_constants)},
+ init = {_optional_source(spec.init)},
+ mutate = {_optional_source(spec.mutate)},
+ is_valid = {_optional_source(spec.is_valid)},
+ init_invalid = {_optional_source(spec.init_invalid)},
+ string = {_optional_source(spec.string)},
+ preamble = {_optional_source(spec.preamble)},
+ definitions = {_optional_source(spec.definitions)},
+ optimizable = {str(spec.can_optimize).lower()},
+ )
+ """)
+ body = _TYPE_SPEC_MODULE.replace("__TYPE_SPEC_CONFIG__", config, 1)
+ fingerprint = hashlib.sha256(body.encode()).hexdigest()
+ module_name = f"_PySRTypeSpec_{fingerprint[:20]}"
+ source = _block(f"""
+ module {module_name}
+ using Random
+ using SymbolicRegression
+ using PythonCall
+
+ {body}
+
+ const _PYSR_PARENT_BINDING_NAMES = Tuple(sort!(
+ filter(
+ name -> isdefined(@__MODULE__, name) &&
+ !startswith(String(name), '#'),
+ collect(names(@__MODULE__; all=true, imported=true)),
+ );
+ by=String,
+ ))
+ const {_TYPE_MODULE_INSTALLED} = true
+ end
+ import .{module_name}: {spec.name}
+ """) + "\n"
+ return _TypeSpecDefinition(
+ spec=copy.deepcopy(spec),
+ fingerprint=fingerprint,
+ source=source,
+ )
+
+
+def load_type_spec_runtime(
+ definition: _TypeSpecDefinition | _TypeSpecRuntimeDefinition,
+) -> _TypeSpecRuntime:
+ """Load deterministic TypeSpec modules and return their Julia objects."""
+ type_definition = (
+ definition.type_spec
+ if isinstance(definition, _TypeSpecRuntimeDefinition)
+ else definition
+ )
+ if isinstance(definition, _TypeSpecRuntimeDefinition):
+ install_source = definition.source
+ filename = "PySR." + definition.module_name
+ else:
+ install_source = type_definition.source
+ filename = "PySR." + type_definition.module_name
+ module_symbol = jl.Symbol(type_definition.module_name)
+ runtime_module_symbol = (
+ jl.Symbol(definition.module_name)
+ if isinstance(definition, _TypeSpecRuntimeDefinition)
+ else None
+ )
+ needs_install = not bool(jl.isdefined(jl.Main, module_symbol))
+ if not needs_install:
+ module = jl.getproperty(jl.Main, module_symbol)
+ needs_install = not bool(
+ jl.isdefined(module, jl.Symbol(_TYPE_MODULE_INSTALLED))
+ )
+ if not needs_install and runtime_module_symbol is not None:
+ needs_install = not bool(jl.isdefined(module, runtime_module_symbol))
+ if not needs_install:
+ configuration_module = jl.getproperty(module, runtime_module_symbol)
+ needs_install = not bool(
+ jl.isdefined(configuration_module, jl.Symbol("_definition_values"))
+ )
+ if needs_install:
+ jl.Base.include_string(jl.Main, install_source, filename)
+ module = jl.getproperty(jl.Main, module_symbol)
+ value_type = jl.getproperty(module, jl.Symbol(type_definition.spec.name))
+
+ if isinstance(definition, _TypeSpecDefinition):
+ runtime = _TypeSpecRuntime(
+ definition=definition,
+ module=module,
+ value_type=value_type,
+ configuration_module=module,
+ operator_names={},
+ operator_functions={},
+ )
+ else:
+ assert runtime_module_symbol is not None
+ configuration_module = jl.getproperty(module, runtime_module_symbol)
+ raw_values = list(
+ jl.getproperty(configuration_module, jl.Symbol("_definition_values"))
+ )
+ source_count = len(
+ _runtime_sources(
+ definition.operators,
+ definition.elementwise_loss,
+ definition.loss_function,
+ definition.loss_function_expression,
+ definition.complexity_mapping,
+ definition.early_stop_condition,
+ definition.expression_spec,
+ )
+ )
+ expected_count = source_count + (
+ definition.expression_spec_function_selector is not None
+ )
+ if len(raw_values) != expected_count:
+ raise RuntimeError(
+ "TypeSpec runtime source/value ordering is inconsistent."
+ )
+ values = iter(raw_values[:source_count])
+ operator_functions = {
+ arity: tuple(next(values) for _ in operator_sources)
+ for arity, operator_sources in definition.operators
+ }
+ optional_sources = (
+ definition.elementwise_loss,
+ definition.loss_function,
+ definition.loss_function_expression,
+ definition.complexity_mapping,
+ definition.early_stop_condition,
+ )
+ optional_values = [
+ next(values) if source is not None else None for source in optional_sources
+ ]
+ runtime = _TypeSpecRuntime(
+ definition=definition,
+ module=module,
+ value_type=value_type,
+ configuration_module=configuration_module,
+ operator_functions=operator_functions,
+ operator_names={
+ arity: tuple(str(jl.Base.nameof(function)) for function in functions)
+ for arity, functions in operator_functions.items()
+ },
+ elementwise_loss=optional_values[0],
+ loss_function=optional_values[1],
+ loss_function_expression=optional_values[2],
+ complexity_mapping=optional_values[3],
+ early_stop_condition=optional_values[4],
+ expression_spec=next(values),
+ )
+ return runtime
+
+
+@cache
+def _type_spec_worker_addprocs_factory() -> AnyValue:
+ return jl.seval(r"""
+ begin
+ import Distributed
+ function (
+ addprocs_function,
+ source,
+ filename,
+ worker_imports,
+ type_module_name,
+ runtime_module_name,
+ )
+ addprocs = something(addprocs_function, Distributed.addprocs)
+ expression = Meta.parseall(source)
+ type_module_symbol = Symbol(type_module_name)
+ runtime_module_symbol = Symbol(runtime_module_name)
+ return function (numprocs; kwargs...)
+ procs = addprocs(numprocs; kwargs...)
+ try
+ SymbolicRegression.import_module_on_workers(
+ procs,
+ pathof(SymbolicRegression),
+ worker_imports,
+ 0,
+ )
+ Distributed.remotecall_eval(Main, procs, expression)
+ head_module = getproperty(
+ getproperty(Main, type_module_symbol),
+ runtime_module_symbol,
+ )
+ module_expression = :(
+ getproperty(
+ getproperty(
+ Main,
+ $(QuoteNode(type_module_symbol)),
+ ),
+ $(QuoteNode(runtime_module_symbol)),
+ )
+ )
+ for proc in procs
+ worker_module = Distributed.remotecall_fetch(
+ Core.eval,
+ proc,
+ Main,
+ module_expression,
+ )
+ worker_module === head_module || error(
+ "TypeSpec runtime identity differs on worker $proc."
+ )
+ end
+ catch
+ Distributed.rmprocs(procs)
+ rethrow()
+ end
+ return procs
+ end
+ end
+ end
+ """)
+
+
+def create_type_spec_addprocs_function(
+ definition: _TypeSpecRuntimeDefinition,
+ addprocs_function: AnyValue | None,
+ worker_imports: AnyValue | None,
+) -> AnyValue:
+ """Wrap worker creation so every new process installs the exact runtime source."""
+ return _type_spec_worker_addprocs_factory()(
+ addprocs_function,
+ definition.source,
+ "PySR." + definition.module_name,
+ worker_imports,
+ definition.type_spec.module_name,
+ definition.module_name,
+ )
+
+
+@cache
+def _type_spec_validator() -> AnyValue:
+ return jl.seval(r"""
+ function (module_, T, type_name, optimizable)
+ DE = SymbolicRegression.InterfaceDynamicExpressionsModule.DE
+ fail(hook, message) = throw(ArgumentError("TypeSpec `$hook` $message"))
+ function call(hook, f, args...)
+ try
+ return f(args...)
+ catch error
+ fail(hook, "failed: $(sprint(showerror, error))")
+ end
+ end
+ function check_value(hook, value)
+ value isa T || fail(hook, "must return `$type_name`.")
+ valid = call("is_valid", DE.is_valid, value)
+ valid isa Bool || fail("is_valid", "must return `Bool`.")
+ valid || fail(hook, "returned an invalid value.")
+ return value
+ end
+ function scalar_constant_count(value)
+ count = call(
+ "count_scalar_constants", DE.count_scalar_constants, value
+ )
+ count isa Int && count >= 0 ||
+ fail("count_scalar_constants", "must return a nonnegative `Int`.")
+ return count
+ end
+ function check_optimization(value, count)
+ scalar_constants =
+ call("scalar_constants", module_._scalar_constants, value)
+ scalar_constants isa AbstractVector ||
+ fail("scalar_constants", "must return an `AbstractVector`.")
+ length(scalar_constants) == count ||
+ fail("count_scalar_constants", "disagrees with `scalar_constants`.")
+ number_type = DE.get_number_type(T)
+ isconcretetype(number_type) && number_type <: AbstractFloat ||
+ fail("scalar_constants", "must return a vector with a concrete `AbstractFloat` element type.")
+
+ offset = 3
+ packed = fill(number_type(NaN), count + 4)
+ next_idx = call(
+ "pack_scalar_constants!",
+ DE.pack_scalar_constants!,
+ packed,
+ offset,
+ value,
+ )
+ next_idx == offset + count ||
+ fail("pack_scalar_constants!", "returned the wrong next index.")
+ all(isnan, packed[1:(offset - 1)]) &&
+ all(isnan, packed[(offset + count):end]) ||
+ fail("pack_scalar_constants!", "wrote outside its scalar-constant range.")
+ isequal(packed[offset:(offset + count - 1)], scalar_constants) ||
+ fail("pack_scalar_constants!", "disagrees with `scalar_constants`.")
+
+ rebuilt = call(
+ "with_scalar_constants",
+ module_._with_scalar_constants,
+ value,
+ scalar_constants,
+ )
+ rebuilt isa T ||
+ fail("with_scalar_constants", "must return `$type_name`.")
+ isequal(rebuilt, value) ||
+ fail("with_scalar_constants(value, scalar_constants(value))", "must preserve the value.")
+
+ result = call(
+ "unpack_scalar_constants",
+ DE.unpack_scalar_constants,
+ packed,
+ offset,
+ value,
+ )
+ result isa Tuple && length(result) == 2 ||
+ fail("unpack_scalar_constants", "must return `(next_idx, $type_name)`.")
+ unpacked_idx, unpacked = result
+ unpacked_idx == offset + count ||
+ fail("unpack_scalar_constants", "returned the wrong next index.")
+ unpacked isa T ||
+ fail("unpack_scalar_constants", "must return `$type_name`.")
+
+ repacked = fill(number_type(NaN), count + 4)
+ repacked_idx = call(
+ "pack_scalar_constants!",
+ DE.pack_scalar_constants!,
+ repacked,
+ offset,
+ unpacked,
+ )
+ repacked_idx == offset + count && isequal(packed, repacked) ||
+ fail("scalar-constant optimization hooks", "must preserve the packed scalar representation.")
+ end
+
+ rng = module_.Random.Xoshiro(0)
+ sampled = check_value("sample", call("sample", SymbolicRegression.sample_value, rng, T, nothing))
+ initial = check_value("init", call("init", SymbolicRegression.init_value, T))
+ for value in (initial, sampled)
+ count = scalar_constant_count(value)
+ optimizable && check_optimization(value, count)
+ end
+
+ if module_._config.init_invalid !== nothing
+ invalid = call("init_invalid", DE.invalid_value, T)
+ invalid isa T || fail("init_invalid", "must return `$type_name`.")
+ valid = call("is_valid", DE.is_valid, invalid)
+ valid isa Bool || fail("is_valid", "must return `Bool`.")
+ !valid || fail("init_invalid", "must return an invalid value.")
+ end
+
+ mutated = check_value(
+ "mutate",
+ call("mutate", SymbolicRegression.mutate_value, rng, sampled, 1.0, SymbolicRegression.ConstantMutation()),
+ )
+ count = scalar_constant_count(mutated)
+ optimizable && check_optimization(mutated, count)
+
+ call("string", module_._string, sampled) isa AbstractString ||
+ fail("string", "must return an `AbstractString`.")
+ return sampled
+ end
+ """)
+
+
+def validate_type_spec_runtime(runtime: _TypeSpecRuntime) -> None:
+ """Invoke TypeSpec hooks once before searching a new configuration."""
+ try:
+ sampled = _type_spec_validator()(
+ runtime.module,
+ runtime.value_type,
+ runtime.spec.name,
+ runtime.spec.can_optimize,
+ )
+ except JuliaError as error:
+ raise ValueError(str(jl.sprint(jl.showerror, error.args[0]))) from error
+ validate_type_spec_options(
+ runtime,
+ runtime.operator_functions,
+ runtime.elementwise_loss,
+ probe_value=sampled,
+ )
+
+
+@cache
+def _type_spec_operator_validator() -> AnyValue:
+ return jl.seval(r"""
+ function (T, type_name, operator, arity, probe_value)
+ inferred = Base.promote_op(operator, ntuple(_ -> T, arity)...)
+ inferred === T || throw(ArgumentError(
+ "TypeSpec operator `$(nameof(operator))` must be type-stable and " *
+ "infer `$type_name` as its return type; inferred `$inferred`. " *
+ "Add an explicit `::$type_name` return annotation if needed."
+ ))
+ if probe_value !== nothing
+ args = ntuple(_ -> probe_value, arity)
+ result = Base.invokelatest(operator, args...)
+ result isa T || throw(ArgumentError(
+ "TypeSpec operator `$(nameof(operator))` must return `$type_name`."
+ ))
+ end
+ return nothing
+ end
+ """)
+
+
+@cache
+def _type_spec_loss_validator() -> AnyValue:
+ return jl.seval(r"""
+ function (T, elementwise_loss, configured_loss_type, probe_value)
+ loss_type = if elementwise_loss === nothing
+ configured_loss_type
+ else
+ Base.promote_op(elementwise_loss, T, T)
+ end
+ isconcretetype(loss_type) && loss_type <: AbstractFloat ||
+ throw(ArgumentError(
+ "The TypeSpec loss must return a concrete subtype of " *
+ "`AbstractFloat`; got `$loss_type`. Add a concrete Julia " *
+ "return type annotation."
+ ))
+ if probe_value !== nothing && elementwise_loss !== nothing
+ loss = Base.invokelatest(
+ elementwise_loss, probe_value, probe_value
+ )
+ loss isa loss_type || throw(ArgumentError(
+ "The TypeSpec elementwise loss must return `$loss_type`."
+ ))
+ end
+ return loss_type
+ end
+ """)
+
+
+def validate_type_spec_options(
+ runtime: _TypeSpecRuntime,
+ operators: dict[int, tuple[AnyValue, ...]],
+ elementwise_loss: AnyValue | None,
+ *,
+ probe_value: AnyValue | None = None,
+) -> AnyValue:
+ """Validate ordinary Julia options against a loaded TypeSpec."""
+ try:
+ for arity, functions in operators.items():
+ for function in functions:
+ _type_spec_operator_validator()(
+ runtime.value_type,
+ runtime.spec.name,
+ function,
+ arity,
+ probe_value,
+ )
+ configured_loss_type = (
+ None
+ if runtime.spec.loss_type is None
+ else jl.Base.include_string(
+ runtime.module,
+ runtime.spec.loss_type,
+ "PySR TypeSpec loss_type",
+ )
+ )
+ return _type_spec_loss_validator()(
+ runtime.value_type,
+ elementwise_loss,
+ configured_loss_type,
+ probe_value,
+ )
+ except JuliaError as error:
+ raise ValueError(str(jl.sprint(jl.showerror, error.args[0]))) from error
+
+
+def type_spec_to_julia_array(
+ runtime: _TypeSpecRuntime, values: Any, *, transpose: bool = False
+) -> AnyValue:
+ """Convert logical Python values to an array of the generated Julia type."""
+ array = (
+ values if isinstance(values, np.ndarray) else np.asarray(values, dtype=object)
+ )
+ if array.dtype != object:
+ array = array.astype(object)
+ if transpose:
+ array = array.T
+ if array.ndim not in (1, 2):
+ raise ValueError("TypeSpec data must be a 1D or 2D array.")
+
+ try:
+ return runtime.module._convert_array(array.ravel(order="F"), array.shape)
+ except JuliaError as error:
+ raise ValueError(str(jl.sprint(jl.showerror, error.args[0]))) from error
+
+
+def type_spec_to_python_array(runtime: _TypeSpecRuntime, values: Any) -> np.ndarray:
+ """Unwrap generated values into their logical one-field or tuple payloads."""
+ field_names = tuple(runtime.spec.fields)
+ output = np.empty(len(values), dtype=object)
+ for i, value in enumerate(values):
+ if len(field_names) == 1:
+ output[i] = getattr(value, field_names[0])
+ else:
+ output[i] = tuple(getattr(value, name) for name in field_names)
+ return output
+
+
+class CallableJuliaExpression:
+ def __init__(self, expression: AnyValue, runtime: _TypeSpecRuntime):
+ self.expression = expression
+ self.runtime = runtime
+
+ def __call__(self, X: np.ndarray, *args):
+ jl_X = type_spec_to_julia_array(
+ self.runtime,
+ object_array_2d(X),
+ transpose=True,
+ )
+ raw_output, completed = SymbolicRegression.eval_tree_array(
+ self.expression, jl_X, *args
+ )
+ if not bool(completed):
+ raise ValueError(
+ "The expression could not be evaluated over the given input: an "
+ "operator returned a value that the spec's `is_valid` rejected."
+ )
+ return type_spec_to_python_array(self.runtime, raw_output)
+
+
+def create_type_spec_exports(
+ runtime: _TypeSpecRuntime,
+ equations: pd.DataFrame,
+ search_output: tuple[AnyValue, AnyValue],
+ output_index: int | None,
+) -> pd.DataFrame:
+ equations = copy.deepcopy(equations)
+ _, all_out_hof = search_output
+ out_hof = all_out_hof[output_index] if output_index is not None else all_out_hof
+ expressions = []
+ callables = []
+ for _, row in equations.iterrows():
+ expression = out_hof.members[row["complexity"] - 1].tree
+ expressions.append(expression)
+ callables.append(CallableJuliaExpression(expression, runtime))
+ return pd.DataFrame(
+ {"julia_expression": expressions, "lambda_format": callables},
+ index=equations.index,
+ )
+
+
+def prepare_type_spec_fit_data(
+ model: Any,
+ X: Any,
+ y: Any,
+ Xresampled: Any,
+ weights: Any,
+ variable_names: Any,
+ complexity_of_variables: Any,
+ X_units: Any,
+ y_units: Any,
+) -> tuple[np.ndarray, np.ndarray, None, Any, np.ndarray, Any, Any, Any]:
+ if Xresampled is not None or model.denoise or model.select_k_features:
+ raise NotImplementedError(
+ "TypeSpec does not support resampling, denoising, or feature selection."
+ )
+ if isinstance(X, pd.DataFrame):
+ if variable_names is not None:
+ warnings.warn(
+ "`variable_names` has been reset to `None` as `X` is a DataFrame."
+ )
+ variable_names = X.columns.astype(str).to_numpy()
+ X = X.to_numpy(dtype=object)
+ else:
+ X = object_array_2d(X)
+ if X.ndim != 2:
+ raise ValueError("TypeSpec X must be a 2D array of logical values.")
+ if X.shape[1] == 0:
+ raise ValueError("TypeSpec X must contain at least one feature.")
+ if variable_names is not None and len(variable_names) != X.shape[1]:
+ raise ValueError("`variable_names` must provide one name per TypeSpec feature.")
+ if variable_names is not None and any(" " in name for name in variable_names):
+ variable_names = [name.replace(" ", "_") for name in variable_names]
+ warnings.warn(
+ "Spaces in variable names are not supported. "
+ "Spaces have been replaced with underscores. \n"
+ "Please use valid names instead."
+ )
+
+ if isinstance(y, (pd.Series, pd.DataFrame)):
+ y = y.to_numpy(dtype=object)
+ else:
+ y = object_array_1d(y)
+ if y.ndim == 2 and y.shape[1] == 1:
+ y = y[:, 0]
+ if y.ndim != 1:
+ raise NotImplementedError("TypeSpec currently supports one output.")
+ if X.shape[0] != y.shape[0]:
+ raise ValueError("X and y have inconsistent numbers of samples.")
+ if X.shape[0] == 0:
+ raise ValueError("X and y must contain at least one sample.")
+ if weights is not None:
+ raise NotImplementedError("TypeSpec does not currently support weights.")
+ if X_units is not None or y_units is not None:
+ raise NotImplementedError("TypeSpec does not currently support units.")
+
+ model.n_features_in_ = X.shape[1]
+ if variable_names is None:
+ variable_names = np.array([f"x{i}" for i in range(X.shape[1])])
+ model.feature_names_in_ = np.asarray(variable_names, dtype=str)
+ model.display_feature_names_in_ = model.feature_names_in_
+ model.nout_ = 1
+ model.complexity_of_variables_ = copy.deepcopy(complexity_of_variables)
+ model.X_units_ = copy.deepcopy(X_units)
+ model.y_units_ = copy.deepcopy(y_units)
+ return (
+ X,
+ y,
+ None,
+ weights,
+ model.feature_names_in_,
+ complexity_of_variables,
+ X_units,
+ y_units,
+ )
+
+
+def prepare_type_spec_prediction_data(model: Any, X: Any) -> np.ndarray:
+ if not isinstance(X, pd.DataFrame):
+ X = object_array_2d(X)
+ if X.ndim != 2:
+ raise ValueError("X must be a 2D array.")
+ if X.shape[1] != model.n_features_in_:
+ raise ValueError("X has a different number of features than during fit.")
+ X = pd.DataFrame(X)
+ if isinstance(X.columns, pd.RangeIndex):
+ if model.selection_mask_ is not None:
+ X = X[X.columns[model.selection_mask_]]
+ X.columns = model.feature_names_in_
+ columns = X.columns.astype(str)
+ if columns.str.contains(" ").any():
+ X = X.copy()
+ X.columns = columns.str.replace(" ", "_")
+ warnings.warn(
+ "Spaces in DataFrame column names are not supported. "
+ "Spaces have been replaced with underscores. \n"
+ "Please rename the columns to valid names."
+ )
+ X = X.rename(columns=str)
+ missing_features = set(model.feature_names_in_) - set(X.columns)
+ if missing_features:
+ raise ValueError(f"X is missing features: {sorted(missing_features)}")
+ return np.asarray(X.reindex(columns=model.feature_names_in_).to_numpy(dtype=object))
+
+
+def validate_type_spec_model_configuration(model: Any) -> None:
+ if model.binary_operators is not None or model.unary_operators is not None:
+ raise ValueError(
+ "TypeSpec requires `operators={...}` and does not accept "
+ "`binary_operators` or `unary_operators`."
+ )
+ validate_type_spec_configuration(
+ model.type_spec,
+ model.operators,
+ elementwise_loss=model.elementwise_loss,
+ loss_function=model.loss_function,
+ loss_function_expression=model.loss_function_expression,
+ )
+ unsupported = {
+ "turbo": model.turbo,
+ "bumper": model.bumper,
+ "autodiff_backend": model.autodiff_backend is not None,
+ "output_jax_format": model.output_jax_format,
+ "output_torch_format": model.output_torch_format,
+ "extra_sympy_mappings": model.extra_sympy_mappings is not None,
+ "extra_jax_mappings": model.extra_jax_mappings is not None,
+ "extra_torch_mappings": model.extra_torch_mappings is not None,
+ }
+ configured = [name for name, enabled in unsupported.items() if enabled]
+ if configured:
+ raise ValueError(
+ "TypeSpec does not support "
+ + ", ".join(f"`{name}`" for name in configured)
+ + "."
+ )
+ expression_spec = model.expression_spec_
+ if not expression_spec.supports_type_spec:
+ raise ValueError(f"{type(expression_spec).__name__} does not support TypeSpec.")
+ expression_spec._validate_type_spec()
diff --git a/release-please-config.json b/release-please-config.json
index 870fa18da..bb49de6be 100644
--- a/release-please-config.json
+++ b/release-please-config.json
@@ -4,8 +4,8 @@
".": {
"release-type": "python",
"include-component-in-tag": false,
- "versioning": "prerelease",
- "prerelease": true
+ "versioning": "default",
+ "prerelease": false
}
}
}
diff --git a/scripts/paper_image_sync.py b/scripts/paper_image_sync.py
index 486ea6926..b037dbf10 100644
--- a/scripts/paper_image_sync.py
+++ b/scripts/paper_image_sync.py
@@ -551,7 +551,7 @@ def main() -> None:
base=docs_default_branch,
title=f"Add paper images from PySR PR #{pr_number}",
body=(
- f"Automated upload of paper image(s) from MilesCranmer/PySR PR #{pr_number}.\n\n"
+ f"Automated upload of paper image(s) from astroautomata/PySR PR #{pr_number}.\n\n"
f"Source head: {pr.head_repo}@{pr.head_sha}\n"
f"Target dir: `{docs_images_dir}`\n"
),
diff --git a/skills/pysr/SKILL.md b/skills/pysr/SKILL.md
new file mode 100644
index 000000000..ac34153f4
--- /dev/null
+++ b/skills/pysr/SKILL.md
@@ -0,0 +1,366 @@
+---
+name: pysr
+description: Use when fitting equations to data with PySR or SymbolicRegression.jl, when a user wants an interpretable formula, symbolic model, scaling law, or empirical relation discovered from numeric data, or when debugging a PySR search that is slow, stuck, or giving poor equations.
+---
+
+# Using PySR Effectively
+
+PySR evolves expression trees through the Julia backend SymbolicRegression.jl, returning readable equations on an accuracy/complexity Pareto front. This guide distills the documentation and several hundred real user threads. Full docs: https://pysr.ai/
+
+## Quick start
+
+```python
+# pip install pysr (Julia is downloaded automatically on first import; no separate install)
+import numpy as np
+from pysr import PySRRegressor
+
+X = 2 * np.random.randn(100, 5)
+y = 2 * np.cos(X[:, 3]) + X[:, 0] ** 2 - 2
+
+model = PySRRegressor(
+ binary_operators=["+", "-", "*", "/"],
+ unary_operators=["cos"],
+ niterations=100,
+)
+model.fit(X, y)
+print(model) # Pareto front: complexity, loss, equation
+model.predict(X) # uses the auto-selected "best" equation
+model.sympy() # SymPy expression
+model.latex() # LaTeX string
+```
+
+`model.equations_` is a pandas DataFrame containing `complexity`, `loss`, `score`, `equation`, `sympy_format`, `lambda_format`. Select any front row by index in `predict`, `sympy`, `latex`, `jax`, `pytorch`; without an index, methods use the auto-selected equation. Indexing lets you inspect and export a specific Pareto candidate instead of only the default row.
+
+## Keep the process alive
+
+The first `.fit()` pays Julia startup/JIT, roughly 1-2 minutes; subsequent same-process fits start almost instantly. Repeated fresh processes repay startup and compilation for every experiment, a common agent failure.
+
+- Keep a long-lived IPython/Jupyter kernel or REPL. New `PySRRegressor` instances are cheap; Julia runtime state lives per-process.
+- Scripts work. Iterate parameters inside one process instead of repeatedly relaunching `python script.py`. Process reuse is an optimization rather than a requirement.
+- Reuse an environment containing `pysr`; fresh environments re-resolve and precompile Julia packages.
+- Expect “Compiling Julia backend...” initially. Avoid treating normal compilation as a hang.
+
+## Candidate equations with `guesses`
+
+Inject knowledge through modifiable/discardable `guesses` or mandatory template structure. Combine both using component/parameter guesses. Constructor strings require Julia syntax, input variable names, and enabled operators:
+
+```python
+model = PySRRegressor(
+ operators={2: ["+", "*"]},
+ guesses=["x0 + 0.5 * x1", "x0 * (x1 + 1.0)"],
+)
+```
+
+Single output: candidate list. Multiple outputs: lists in target-column order, `guesses=[["x0 + x1"], ["x0 * x1"]]`. Templates: dictionaries keyed by component/parameter names; `#1`, `#2`, etc. identify component arguments.
+
+Guesses enter populations throughout search. `should_optimize_constants=True` optimizes constants/template parameters before insertion. `fraction_replaced_guesses`, default `0.001`, controls the cycle-end replacement fraction.
+
+Guesses are pushed into the hall of fame before generation one. Inspect accepted seeds at their complexity in `model.equations_` or immediate `hall_of_fame.csv`; confirm insertion and optimized constants. Oversize seeds produce only Julia warning `Guess expression '...' has complexity 23 > maxsize (10)`, using default `x1, x2, ...` names regardless of supplied names. They remain unrecorded but still migrate; descendants become recordable after mutations satisfy `maxsize`. Violations of `constraints`, `nested_constraints`, or `maxdepth` disappear silently. Lower-loss candidates replace seeds at the same complexity: only one expression survives per complexity.
+
+Keep constraints compatible and set `maxsize` slightly above the largest seed; use the warning's complexity. Cost rises steeply, PySR warns above 40, and several-times-seed sizes can exhaust the budget without completing one generation despite appearing alive. Use `warmup_maxsize_by` against early bloat.
+
+Set guesses on the estimator, never `.fit()`. Continue with `model.set_params(guesses=[...], warm_start=True)`, replacing guesses while keeping data representation/search space fixed.
+
+## Recommended workflow
+
+1. **Subsample:** usually 1,000-5,000 representative rows suffice even from millions. Many features, heavy noise, or rare regimes benefit from more. PySR 2.0 automatically batches larger datasets; `batching=False` forces full-data search evaluations. Batching still reevaluates hall-of-fame candidates fully. Fewer rows proportionally accelerate search.
+2. **Minimize operators:** select domain-plausible operators. Redundant `pow`/`square`/`cube` or `-`/`neg` enlarge search. Polynomials: `["+", "-", "*"]`, omit `/` and `^`.
+3. **Use defaults:** populations, parsimony, mutation weights, and `ncycles_per_iteration` were tuned by large-scale search in 2024-2025. Avoid older forum/paper recipes.
+4. **Debug briefly:** check operators, loss, and sensible equations in a few-minute run before one long search.
+5. **Budget the final run:** evolution lacks conventional convergence and can jump families after hours. Set large `niterations`; control wall-clock with `timeout_in_seconds`, total evaluations with `max_evals`, or stopping with `early_stop_condition="stop_if(loss, complexity) = loss < 1e-6 && complexity < 10"` (Julia `(loss, complexity)` function). `max_evals` supports compute-matched comparisons; `niterations * populations * population_size` does not count evaluations. Monitor partial `outputs//hall_of_fame.csv`; stop gracefully with `q` then Enter in IPython.
+6. **Inspect the entire front:** validate candidates and report tradeoffs.
+
+## Choosing an equation
+
+`model_selection="best"` maximizes `score` within 1.5x minimum loss; `"accuracy"` minimizes loss. Score is negative log-loss slope per complexity unit; a large score marks a sharp accuracy gain for little added complexity. These heuristics control default `predict`.
+
+Print `model.equations_`; inspect sharp loss drops, held-out `model.predict(X_test, index=i)` when overfitting is plausible, and limits x -> 0/inf against domain expectations. Present 2-3 candidates with tradeoffs.
+
+## Losses and weights
+
+Default: MSE.
+
+- Outliers: `elementwise_loss="L1DistLoss()"`, median-seeking versus mean-seeking.
+- Uncertainty: `model.fit(X, y, weights=1/sigma**2)`; built-ins apply weights automatically. Custom: `elementwise_loss="myloss(x, y, w) = w * abs(x - y)^2"`.
+- Wide target range: largest values dominate MSE; use log-space loss:
+
+```python
+elementwise_loss = """function loss_fnc(prediction, target)
+ scatter_loss = abs(log((abs(prediction)+1e-20) / (abs(target)+1e-20)))
+ sign_loss = 10 * (sign(prediction) - sign(target))^2
+ return scatter_loss + sign_loss
+end"""
+```
+
+- Relative/percentage error: divide by target; prediction denominators reward infinite predictions. Guard near-zero targets.
+- LossFunctions.jl strings work: `"HuberLoss(1.0)"`, `"LPDistLoss{3}()"`.
+- Binary classification: +1/-1 targets, margin loss `"L2MarginLoss()"`; apply sigmoid yourself for probabilities.
+- Asymptotes/boundaries, including y -> 0 as x -> inf, exact boundary values, and known limits: add synthetic points with very large `weights`, usually preferable to custom objectives. Strict enforcement: estimate limits numerically inside `loss_function`, add graded penalties.
+- Keep losses deterministic for caching and non-negative; `loss_scale="linear"` permits negatives, including log-likelihoods.
+
+`elementwise_loss` receives scalar `(prediction, target)` or `(prediction, target, weight)`; never sum/broadcast. Whole-vector/tree objectives require `loss_function`.
+
+## Scaling and features
+
+- Normalization is optional: constants sample near N(0,1), mutate multiplicatively; extreme scales slow search. Prefer natural units, avoiding nuisance constants and hidden physical meaning. Rescale visibly struggling searches.
+- Roughly 10 features or fewer need no special handling.
+- Tens: increase `maxsize` and rows; search selects features reasonably well, automatic batching handles volume. Requiring all 30+ features needs `maxsize` well above 100.
+- Above ~50: engineer domain aggregates or template decompositions. Fields/images/sequences/graphs need meaningful features rather than pixel columns, or appropriately biased neural networks followed by symbolic component distillation, arXiv:2006.11287. Fallback: gradient-boosting preselection, `select_k_features=k`.
+- Omitted variables failed to justify complexity through accuracy; force inclusion through an absence-penalizing custom loss.
+
+## PDE discovery
+Regress `u_t = f(u, u_x, u_xx, ...)` from `u(x, t)`: grid points become samples, field/spatial derivatives become features, time derivative becomes target; PySR discovers products like `u*u_x` without a fixed library.
+- Spatial derivatives: spectral differentiation for clean periodic grids, otherwise Savitzky-Golay, `savgol_filter(U, 21, 3, deriv=1, delta=dx, axis=-1)`; finite differences require clean data, spectral derivatives catastrophically amplify noise.
+- Target: central snapshot differences; smooth noisy/sparse time samples along time or sample finer. Trim non-periodic stencil margins.
+- Bias low-order terms with column-ordered `complexity_of_variables=[1, 2, 3, ...]`, costing `u`, `u_x`, `u_xx` as 1, 2, 3; prune dimensional impossibilities with `X_units=["m/s", "s^-1", ...]`. Inspect the front elbow; noisy coefficients are approximate. Validate forward simulation from held-out initial conditions.
+## Template expressions
+Default to plain search; suggest `TemplateExpressionSpec` for known structure that free-form trees would rediscover or violate:
+- Outer forms: `sin(f(x1, x2)) + g(x3)`, rational functions, envelopes `x*(1-x)*g(x)`.
+- Per-category/class/object/condition coefficients through `parameters`; shared subexpressions/coupled outputs; derivative/integral relations through differential operator `D`; hard variable-placement requirements.
+```python
+from pysr import PySRRegressor, TemplateExpressionSpec
+spec = TemplateExpressionSpec(
+ combine="sin(f(x1, x2)) + g(x3)^2",
+ expressions=["f", "g"],
+ variable_names=["x1", "x2", "x3"],
+)
+model = PySRRegressor(expression_spec=spec, binary_operators=["+", "-", "*", "/"])
+model.fit(X, y)
+```
+Pass category as an X column; increment zero-based IDs for Julia's 1-based indexing:
+```python
+spec = TemplateExpressionSpec(
+ combine="p[class] * f(x1, x2) + q[class]",
+ expressions=["f"],
+ variable_names=["x1", "x2", "class"],
+ parameters={"p": 3, "q": 3},
+)
+model = PySRRegressor(
+ expression_spec=spec,
+ operators={2: ["+", "*"]},
+ guesses=[{
+ "f": "#1 + #2",
+ "p": [5.0, 10.0, 0.8],
+ "q": [0.0, 0.0, 0.0],
+ }],
+)
+```
+`parameters={"p": 3, "q": 3}` declares vector lengths; guesses initialize vectors and `f`. Declare disjoint component/parameter names and every guess key; supply every component formula. Match parameter lengths using Python lists/1D NumPy arrays; omit vectors for normal initialization. Values remain learnable, are copied without modifying caller arrays, and use model precision. Parameter guesses require SymbolicRegression.jl ≥2.3.0. `#1`/`#2` denote passed arguments, here `x1`/`x2`; with `combine="f(x, p[1])"`, `#2` denotes `p[1]`.
+`combine` accepts arbitrary Julia/multiple statements: reuse `fx = f(x); fx + fx^2`, change arguments `f(x1) - f(x2)`, differentiate `df = D(f, 1); df(x)`. Vector/multi-output problems: append extra targets to X, return per-row residuals, fit dummy y with `elementwise_loss="(p, t) -> p"` so output itself is loss.
+PySR 2.0 caveats: arbitrary Julia prevents `.sympy()`, `.latex()`, `.jax()`, `.pytorch()` export; manually reassemble component strings from `model.equations_`. Arguments print positionally as `#1, #2` because components may receive different inputs. Combine values are `ValidVector`s (`.x` data, `.valid` flag); arithmetic propagates validity, custom manipulation must unwrap/rebuild `ValidVector(raw, valid)`. Use Float32-safe `0.5f0` or explicit conversion; bare Float64 `0.5` risks type instability. Custom objectives require `loss_function_expression`. Learnable parameters use `TemplateExpressionSpec(parameters=...)`; pre-1.4 `function_symbols`/lambda-style combine API is removed in 2.0.
+
+## Custom value types: `TypeSpec`
+`TypeSpec` generates Julia structs for nodes, independently of `expression_spec`, including fixed structures. Supply fields in declaration order: `triple`, then `code`; the nested tuple remains one field. Keep X 2D/y 1D and assign object cells individually to prevent extra NumPy axes:
+```python
+import numpy as np
+from pysr import PySRRegressor, TemplateExpressionSpec, TypeSpec
+value = ((1.0, 2.0, 3.0), np.uint32(7))
+X = np.empty((1, 1), dtype=object)
+y = np.empty(1, dtype=object)
+X[0, 0] = value
+y[0] = value
+packet = TypeSpec(
+ "Packet",
+ fields={"triple": "NTuple{3,Float64}", "code": "UInt32"},
+ sample="rng -> Packet(ntuple(_ -> randn(rng), 3), rand(rng, UInt32))",
+ mutate="""function (rng, value::Packet, temperature::Float64)
+ triple = ntuple(i -> value.triple[i] + temperature * randn(rng), 3)
+ bit = UInt32(1) << rand(rng, 0:31)
+ Packet(triple, xor(value.code, bit))
+ end""",
+ scalar_constants="value -> collect(value.triple)",
+ with_scalar_constants="(value, c) -> Packet((c[1], c[2], c[3]), value.code)",
+ is_valid="value -> all(isfinite, value.triple)",
+)
+```
+- `sample(rng) -> Packet`: constant leaves; `mutate(rng, value, temperature) -> Packet`: continuous/discrete mutation; annotations are optional.
+- Supply `scalar_constants(Packet) -> Vector{Float64}` and `with_scalar_constants(Packet, Vector{Float64}) -> Packet` together; expose optimizable scalars, preserve other fields when rebuilding.
+- `is_valid(Packet) -> Bool`: reject invalid intermediates; optional `init`: `() -> Packet`; `string`: `Packet -> AbstractString`.
+- `preamble`: Julia before generated type/hooks; `definitions`: Julia immediately after type, including constructors/type-specific methods.
+- `loss_type`: concrete `AbstractFloat` for full objectives; unavailable for elementwise losses, whose return type is inferred.
+Configure `PySRRegressor(type_spec=packet, operators=operators, elementwise_loss=loss)`. In arity-keyed `operators`, unary `negate_packet(x::Packet)::Packet = Packet(ntuple(i -> -x.triple[i], 3), x.code)` preserves code; binary `combine_packets(x::Packet, y::Packet)::Packet = Packet(ntuple(i -> x.triple[i] + y.triple[i], 3), xor(x.code, y.code))` sums triples/XORs codes. Include binary `choose_parameter(a::Packet, b::Packet) = a` and `choose_parameter(a::Packet, b::ValidVector) = ValidVector(map(_ -> a, b.x), b.valid)`. Define `loss` as `packet_loss(p::Packet, y::Packet)::Float64 = sum((p.triple[i] - y.triple[i])^2 for i in 1:3) + Float64(count_ones(xor(p.code, y.code)))`; annotation keeps return type concrete.
+Typed constructors work in `ordinary_guess = ["combine_packets(x0, Packet((1.0, 2.0, 3.0), UInt32(7)))"]` and `template_component_guess = [{"f": "combine_packets(#1, Packet((1.0, 2.0, 3.0), UInt32(7)))"}]`. Parsing uses the generated configuration module with Packet/operators defined; no parser closure/import needed.
+Template parameters share the data's `TypeSpec`. Define `packet_template = TemplateExpressionSpec(combine="choose_parameter(p[1], f(x))", expressions=["f"], variable_names=["x"], parameters={"p": 1})`. Preserve one logical Packet with `packet_guess_from_list = [{"f": "#1", "p": [value]}]`, or assign `packet_parameter_array = np.empty(1, dtype=object)` and `packet_parameter_array[0] = value`, then `packet_guess_from_array = [{"f": "#1", "p": packet_parameter_array}]`. Configure the estimator above with additional `expression_spec=packet_template`, `guesses=packet_guess_from_array`.
+Verify field count/order, conversion to declared Julia types, type-stable Packet-returning operators, and one real scalar per loss pair; reserve `Inf` for invalid evaluations.
+
+## Custom operators
+
+Supply Julia strings and SymPy mappings using SymPy functions, never numpy/scipy; incorrect mappings break export/prediction:
+
+```python
+model = PySRRegressor(
+ binary_operators=["+", "*"],
+ unary_operators=["inv(x) = 1/x", "gauss(x) = exp(-x^2)"],
+ extra_sympy_mappings={
+ "inv": lambda x: 1/x,
+ "gauss": lambda x: sympy.exp(-x**2),
+ },
+)
+```
+
+- Accept all reals, including probes beyond data. Guard domains: `my_sqrt(x) = x >= 0 ? sqrt(x) : convert(typeof(x), NaN)`; any data-point NaN discards candidates with infinite loss. Prefer protected built-ins `sqrt`, `log`, `acosh`.
+- Preserve types: `T(2.5)` with `where {T}`, or default-Float32 `2.5f0`; bare `2.5`/`0` introduce Float64/Int64 and break Float32 pipelines.
+- One/two scalar inputs, scalar output; ≥3 arguments require arity-keyed `operators` or templates.
+- Julia packages: `from pysr import jl; jl.seval("import Pkg; Pkg.add(\"SpecialFunctions\")"); jl.seval("using SpecialFunctions")`; safely wrap any importable function.
+- No closed SymPy form: map to `class myop(sympy.Function): pass`; symbolic export works, `predict` requires numeric mapping. Call `model.refresh()` after mapping changes.
+
+## Custom objectives: `loss_function`
+
+Use full objectives for vectors, trees, derivatives, auxiliary data:
+
+```python
+objective = """
+function my_objective(tree, dataset::Dataset{T,L}, options) where {T,L}
+ prediction, completed = eval_tree_array(tree, dataset.X, options)
+ !completed && return L(Inf)
+ residuals = prediction .- dataset.y
+ return sum(abs2, residuals) / dataset.n
+end
+"""
+model = PySRRegressor(loss_function=objective, binary_operators=["+", "*", "-"])
+```
+
+- Check `completed` before accessing predictions; failure leaves garbage. Return `L(Inf)` for numerical invalidity; use graded finite structural penalties such as `L(1e6 * n_violations)` so intermediate mutations survive. All-or-nothing structural rejection can make targets unreachable.
+- `dataset.X`: features × samples, transposed from Python; `dataset.y`: vector; `dataset.n`: sample count. Three-argument objectives receive active batches; `(tree, full_dataset, options, idx)` exposes full data/selected indices.
+- Millions of calls: type-stable, vectorized, no printing/Python callbacks (GIL serialization). Use `const` globals or string interpolation, Distances.jl/standard packages; diagnose with `@code_warntype`/BenchmarkTools in separate Julia. Avoid internal-name collision `eval_loss`.
+- Append auxiliary targets/derivatives/group IDs to X and slice, or interpolate literal arrays. Appended columns remain searchable unless templates exclude them or penalties discourage use.
+- Candidate derivatives: `eval_diff_tree_array` for one feature, `eval_grad_tree_array` for all; support monotonicity/physics-informed losses. Templates use `D(f, i)`.
+- Reinterpreted trees, including split subtrees/recursion, print raw equations requiring manual decoding; symbolic manipulation generally breaks `.sympy()`/`.predict()`, requiring manual extraction.
+
+### Structural tree constraints
+
+Before writing a traversal, check whether the rule is already an option. Operator placement rules are `nested_constraints`: "no sin inside cos and no cos inside sin" is `nested_constraints={"sin": {"cos": 0}, "cos": {"sin": 0}}`; "exponent must be a bare constant or variable" is `constraints={"^": (-1, 1)}` (keys are the names used in the operator lists, so `"pow"` if that is the name given). Both reject a violating mutation before it is evaluated, which is cheaper than a loss penalty and never lets a violator into the hall of fame. Use traversal only for what they cannot express: required features, constant value ranges, rules that depend on position or on several operators at once.
+
+Implement those by traversal. Supported: `any(f, tree)`, `all(f, tree)`, `count(f, tree)`, `sum(f, tree)`, `foreach(f, tree)`, `collect(tree)`, depth-first `for node in tree`. Prefer functional traversal: loops allocate stacks; `count`/`sum`/`any` traverse directly.
+
+- `node.degree`: leaf 0/unary 1/binary 2; `node.l`, `node.r`: traversable children.
+- Leaves: `node.constant` distinguishes constants/variables; `node.val`: constant value; `node.feature`: 1-based feature index.
+- `node.op`: 1-based supplied operator-list index, per arity.
+
+Allow `^` only with a lone constant exponent in [0,1]:
+
+```python
+objective = """
+function constrained_loss(tree, dataset::Dataset{T,L}, options) where {T,L}
+ idx_pow = 3 # position of ^ in binary_operators below (1-indexed)
+ n_bad = count(tree) do node
+ node.degree == 2 && node.op == idx_pow &&
+ any(c -> !(c.degree == 0 && c.constant && 0 <= c.val <= 1), node.r)
+ end
+ n_bad > 0 && return L(10_000 * n_bad)
+ prediction, valid = eval_tree_array(tree, dataset.X, options)
+ !valid && return L(Inf)
+ return sum(i -> abs2(prediction[i] - dataset.y[i]), eachindex(prediction)) / dataset.n
+end
+"""
+model = PySRRegressor(binary_operators=["+", "*", "^"], loss_function=objective)
+```
+
+Count violations before evaluation; multiply by a large finite penalty to guide evolution. Require feature 2: `any(n -> n.degree == 0 && !n.constant && n.feature == 2, tree) || return L(big)`. Templates: obtain component trees using `get_tree(ex)` inside `loss_function_expression`. `tree_mapreduce`, `NodeSampler`, construction: https://ai.damtp.cam.ac.uk/dynamicexpressions/stable/examples/base_operations/
+
+## Constraints and complexity
+
+- `constraints={"pow": (9, 1)}` caps argument complexity: base 9, lone constant/variable exponent; `-1` unlimited. Strongly constrain `^`; unrestricted exponentiation searches poorly.
+- `nested_constraints={"sin": {"sin": 0, "cos": 0}, "cos": {"sin": 0, "cos": 0}}` prohibits nested trig; values cap inner occurrences within outer operators.
+- `complexity_of_operators={"exp": 3}`, `complexity_of_constants=2` discourage operators/constants; `complexity_of_variables` accepts global/per-feature costs.
+- Constraints affect every intermediate; excessive tightness silently excludes targets. Allow slack: final size 30, `maxsize=35`.
+- `maxsize` counts operator/constant/variable nodes or `complexity_mapping` units; seven-feature linear models cost ~29, default 30 excludes many-term formulas.
+- `warmup_maxsize_by=0.5` ramps maxsize over that run fraction, helping avoid early complex traps.
+- `parsimony`/`adaptive_parsimony_scaling`: defaults well-tuned in 1.5+.
+- Known pi/G: recognize fitted floats afterwards, or supply exact constant-valued features.
+- Integers: constant features plus increased `complexity_of_constants`, or custom-operator rounding; continuous optimization does not guarantee integers.
+
+### `complexity_mapping`
+
+Replace fixed operator prices/node counts with a Julia expression-to-`Int` function for emitted characters, evaluation cost, tokens, or depth. Pricing applies to front, `maxsize`, parsimony, `constraints`.
+
+```python
+complexity_mapping = """
+function source_chars(ex)
+ tree_mapreduce(
+ leaf -> leaf.constant ? length(string(leaf.val)) : 2, # "x0" is 2 chars
+ branch -> branch.degree == 1 ? 6 : 5, # "sin(" + ")" / " + " with parens
+ (parent, children...) -> parent + sum(children),
+ get_tree(ex),
+ )
+end
+"""
+model = PySRRegressor(
+ binary_operators=["+", "-", "*"],
+ complexity_mapping=complexity_mapping,
+ maxsize=280,
+)
+```
+
+This mapping tracked every test-row `len(equation)` within offset 2 for outer parentheses; rank by actual cost and calibrate against emitted candidates.
+
+- Input: `AbstractExpression`; `get_tree(ex)` exposes `degree`, `constant`, `val`, `feature`, `op`. `tree_mapreduce`/`get_tree` are in scope.
+- Avoid `length(string_tree(ex))`: internal rendering uses `binary_operator[3](...)` and `x1, x2` placeholders, unrelated to exported length.
+- Set real-budget `maxsize`; default 30 with character pricing leaves `equations_` empty when nothing fits. Hall of fame allocates one slot/unit through maxsize; filter long fine-grained fronts.
+- Every candidate incurs computation: cheap, type-stable, no Python callbacks.
+- `guesses` use identical pricing/acceptance rules.
+
+## Dimensional constraints
+
+Check physical units during search:
+
+```python
+model.fit(X, y, X_units=["Constants.M_sun", "kg", "m"], y_units="kg * m / s^2")
+```
+
+- DynamicQuantities.jl notation; `"1"` explicitly dimensionless.
+- Soft `dimensional_constraint_penalty`, default 1000; avoid 1e9, preserve graded paths through imperfect intermediates.
+- Fitted constants have wildcard `[⋅]`/`[?]` units; a lone constant can absorb any dimensions. Forbid with `dimensionless_constants_only=True`.
+- SI radians are dimensionless; units cannot restrict variables to trig.
+
+## Parallelism
+
+- Default `parallelism="multithreading"` suits laptops/single nodes. Thread count is fixed at Julia startup: set `PYTHON_JULIACALL_THREADS=` before importing pysr; `JULIA_NUM_THREADS` is inappropriate under juliacall.
+- Keep `populations` ~2-3x threads/cores for worker utilization; default `populations=31` covers typical machines.
+- Saturated coordinating thread: increase `ncycles_per_iteration`, reducing communication.
+- Long runs: `parallelism="multiprocessing"`, `procs=n` costs substantially more startup/fit, potentially faster steady-state. Multiple nodes: `cluster_manager="slurm"` or native SlurmClusterManager.jl, https://julia.pysr.ai/dev/slurm/. Launch once on one node, let it spawn workers, avoid `srun`; put worker packages in `worker_imports`.
+- Reproducibility: `deterministic=True, random_state=, parallelism="serial"`; parallel seeds remain nondeterministic. CPU differences can persist serially; reduce with `precision=64`.
+
+## Saving, resuming, exporting
+
+- Non-temporary fits write continuously readable `outputs//hall_of_fame.csv` and `checkpoint.pkl`. Reload `PySRRegressor.from_file(run_directory=...)` with identical PySR version; pickles are version-locked. Preserve CSV/construction code durably; resupply custom `extra_sympy_mappings`.
+- `warm_start=True` resumes previous populations on next same-process `.fit()`. Fix operators, `maxsize`, `expression_spec`, precision, feature count/order; loss/weights may change for staged objectives.
+- Export `model.sympy(i)`, `model.latex(i)`, `model.latex_table()`, differentiable `model.jax(i)` returning `{'callable', 'parameters'}`, trainable `model.pytorch(i)`. Supply custom `extra_jax_mappings`/`extra_torch_mappings`; fine-tune selected equations' constants by full-data gradient descent.
+- Julia losses can differ from NumPy recomputation through default 32-bit versus 64-bit precision. Use `precision=64` when important or values exceed ~1e19/fall below ~1e-19, risking Float32 overflow.
+
+## Troubleshooting
+
+| Symptom | Cause / fix |
+|---|---|
+| First import/fit takes minutes | Initial Julia download; per-process JIT. Retain process. |
+| Compilation hangs | Apple Silicon: replace Rosetta x86 Python with native arm64. Notebook/embedded stdin wedge: `input_stream="devnull"`. |
+| Startup `lock.pid` | Killed-process juliapkg lock: verify no installation, delete lock, import once freshly. Shared filesystems: pre-initialize before parallel jobs. |
+| Jupyter `UnicodeDecodeError` | Old PythonCall bug: upgrade or set `PYTHON_JULIACALL_AUTOLOAD_IPYTHON_EXTENSION=no` before import. |
+| Cannot interrupt notebook | Known limitation; IPython/terminal `q`+Enter only. Set `timeout_in_seconds` or use IPython. |
+| Nonsensical equations | Check scale/loss (log-space), missing/redundant operators, insufficient `maxsize`, excluding constraints, in order before more compute. |
+| Tiny equations | Insufficient `maxsize` or MSE dominated by largest points. |
+| `DomainError` | Guard custom operators with typed NaN. |
+| Prediction/SymPy fails | Correct `extra_sympy_mappings`; use SymPy functions. |
+| Division near zero | Any-NaN-anywhere rejects candidates; survivors remain finite only on supplied data, potentially diverging elsewhere. |
+| Early-stop `ProcessExitedException` | Harmless multiprocessing worker-teardown noise. |
+| Growing memory | Mostly fixed in recent Julia; upgrade or set multiprocessing `heap_size_hint_in_bytes`. |
+| Different results | Stochastic evolution; apply serial determinism recipe. |
+
+## Dropping to Julia
+
+Use native SymbolicRegression.jl, `SRRegressor` via MLJ, for Julia pipelines/deep customization: expression types, mutation operators, per-component constraints. `from pysr import jl` exposes the runtime; `jl.seval(...)` executes arbitrary Julia, installed packages support operators/losses. Small readable backend extension points: `src/Options.jl`, `src/CheckConstraints.jl`; connect dev checkouts through `pysr/juliapkg.json`.
+
+## Version notes
+
+Target: PySR 2.0.0, SymbolicRegression.jl 2.x. Relative to 1.5.x: n-ary `operators={1: ["sin"], 2: ["+", "*"], 3: ["clamp"]}`, `guesses`, `TypeSpec`, automatic batching, autodiff plugins. Defaults: `batching="auto"`, `batch_size=None`, `annealing=True`, `crossover_probability=0.2`, new mutation mix; use 2.0 defaults rather than 1.x tuning.
+
+| Deprecated | Current |
+|---|---|
+| `pysr.install()`, `python -m pysr install`, PyCall/`from julia import Main` | Automatic Julia installation on import |
+| `pysr(...)`, `best()`, `get_hof()` | `PySRRegressor`, `.fit()`, `model.equations_`, `get_best()` |
+| `multithreading=True/False`, `procs=0` | `parallelism="multithreading"/"multiprocessing"/"serial"` |
+| `full_objective=` | `loss_function`; templates: `loss_function_expression` |
+| `loss=` | `elementwise_loss` |
+| `equation_file=` | `output_directory=` + `run_id=` |
+| `npop`, `ncyclesperiteration`, camelCase args | `population_size`, `ncycles_per_iteration`, snake_case |
+| Positional `TemplateExpressionSpec`, `function_symbols=...` | Explicit `combine=`, `expressions=`, `variable_names=` |
\n\n### Title\n\n
\ncontent\n...
...\s*
"
- content = re.sub(details_pattern, replace_details, content, flags=re.DOTALL)
-
- output_path.write_text(content)
- print(f"Generated {output_path} from {readme_path}")
-
-
def main():
"""Main entry point."""
# Get the directory containing this script
script_dir = Path(__file__).parent.resolve()
src_dir = script_dir / "src"
- # Generate index.md from README.md
- readme_path = script_dir.parent / "README.md"
- index_path = src_dir / "index.md"
- if readme_path.exists():
- try:
- process_readme_to_index(readme_path, index_path)
- except Exception as e:
- print(f"Error generating index.md: {e}", file=sys.stderr)
- import traceback
-
- traceback.print_exc()
- sys.exit(1)
-
# Process _api.md → api.md and _api-advanced.md → api-advanced.md
for source_name, output_name in [
("_api.md", "api.md"),
diff --git a/docs/gen_vitepress_index.py b/docs/gen_vitepress_index.py
index 3806d12b0..5ddbab121 100644
--- a/docs/gen_vitepress_index.py
+++ b/docs/gen_vitepress_index.py
@@ -3,7 +3,30 @@
Generate VitePress-compatible index.md by combining _index.md template with README.md content.
"""
+import re
from pathlib import Path
+from urllib.parse import urlsplit
+
+REPO_ROOT = Path(__file__).resolve().parent.parent
+REPO_URL = "https://github.com/astroautomata/PySR"
+
+
+def repo_link(match):
+ url = match.group(2)
+ parts = urlsplit(url)
+ if parts.scheme or parts.netloc or url.startswith(("/", "#")):
+ return match.group(0)
+
+ path = (REPO_ROOT / parts.path).resolve()
+ if not path.is_relative_to(REPO_ROOT) or not path.exists():
+ return match.group(0)
+ if path.is_relative_to(REPO_ROOT / "docs" / "src") and path.suffix == ".md":
+ return match.group(0)
+
+ kind = "tree" if path.is_dir() else "blob"
+ suffix = f"?{parts.query}" if parts.query else ""
+ suffix += f"#{parts.fragment}" if parts.fragment else ""
+ return f"{match.group(1)}{REPO_URL}/{kind}/master/{path.relative_to(REPO_ROOT).as_posix()}{suffix})"
def process_readme_content(readme_content):
@@ -24,20 +47,13 @@ def process_readme_content(readme_content):
if not found_first_header:
continue
- # Skip GitHub user-images URLs
- if "user-images" in line:
- continue
-
# Adjust header levels: # -> ##
if line.startswith("# ") and not line.startswith("# PySR"):
line = "#" + line
- # Transform "### Test status" to "**Test status**"
- if line == "### Test status":
- line = "**Test status**"
-
# Fix internal links to use VitePress paths
- line = line.replace("(https://ai.damtp.cam.ac.uk/pysr/papers)", "(/papers)")
+ line = line.replace("(https://pysr.ai/papers)", "(/papers)")
+ line = re.sub(r"(?=14.6"
}
@@ -1394,6 +1397,7 @@
"version": "4.1.3",
"resolved": "https://registry.npmjs.org/ansi-colors/-/ansi-colors-4.1.3.tgz",
"integrity": "sha512-/6w/C21Pm1A7aZitlI5Ni/2J6FFQN8i1Cvz3kHABAAbw93v/NlvKdVOqz7CCWz/3iv/JplRSEEZ83XION15ovw==",
+ "license": "MIT",
"engines": {
"node": ">=6"
}
@@ -1414,7 +1418,8 @@
"node_modules/boolbase": {
"version": "1.0.0",
"resolved": "https://registry.npmjs.org/boolbase/-/boolbase-1.0.0.tgz",
- "integrity": "sha512-JZOSA7Mo9sNGB8+UjSgzdLtokWAky1zbztM3WRLCbZ70/3cTANmQmOdR7y2g+J0e2WXywy1yS468tY+IruqEww=="
+ "integrity": "sha512-JZOSA7Mo9sNGB8+UjSgzdLtokWAky1zbztM3WRLCbZ70/3cTANmQmOdR7y2g+J0e2WXywy1yS468tY+IruqEww==",
+ "license": "ISC"
},
"node_modules/ccount": {
"version": "2.0.1",
@@ -1447,6 +1452,7 @@
"version": "1.0.0-rc.10",
"resolved": "https://registry.npmjs.org/cheerio/-/cheerio-1.0.0-rc.10.tgz",
"integrity": "sha512-g0J0q/O6mW8z5zxQ3A8E8J1hUgp4SMOvEoW/x84OwyHKe/Zccz83PVT4y5Crcr530FV6NgmKI1qvGTKVl9XXVw==",
+ "license": "MIT",
"dependencies": {
"cheerio-select": "^1.5.0",
"dom-serializer": "^1.3.2",
@@ -1467,6 +1473,7 @@
"version": "1.6.0",
"resolved": "https://registry.npmjs.org/cheerio-select/-/cheerio-select-1.6.0.tgz",
"integrity": "sha512-eq0GdBvxVFbqWgmCm7M3XGs1I8oLy/nExUnh6oLqmBditPO9AqQJrkslDpMun/hZ0yyTs8L0m85OHp4ho6Qm9g==",
+ "license": "BSD-2-Clause",
"dependencies": {
"css-select": "^4.3.0",
"css-what": "^6.0.1",
@@ -1491,6 +1498,7 @@
"version": "6.2.1",
"resolved": "https://registry.npmjs.org/commander/-/commander-6.2.1.tgz",
"integrity": "sha512-U7VdrJFnJgo4xjrHpTzu0yrHPGImdsmD95ZlgYSEajAn2JKzDhDTPG9kBTefmObL2w/ngeZnilk+OV9CG3d7UA==",
+ "license": "MIT",
"engines": {
"node": ">= 6"
}
@@ -1510,6 +1518,7 @@
"version": "4.3.0",
"resolved": "https://registry.npmjs.org/css-select/-/css-select-4.3.0.tgz",
"integrity": "sha512-wPpOYtnsVontu2mODhA19JrqWxNsfdatRKd64kmpRbQgh1KtItko5sTnEpPdpSaJszTOhEMlF/RPz28qj4HqhQ==",
+ "license": "BSD-2-Clause",
"dependencies": {
"boolbase": "^1.0.0",
"css-what": "^6.0.1",
@@ -1525,6 +1534,7 @@
"version": "6.2.2",
"resolved": "https://registry.npmjs.org/css-what/-/css-what-6.2.2.tgz",
"integrity": "sha512-u/O3vwbptzhMs3L1fQE82ZSLHQQfto5gyZzwteVIEyeaY5Fc7R4dapF/BvRoSYFeqfBk4m0V1Vafq5Pjv25wvA==",
+ "license": "BSD-2-Clause",
"engines": {
"node": ">= 6"
},
@@ -1538,9 +1548,10 @@
"integrity": "sha512-M1uQkMl8rQK/szD0LNhtqxIPLpimGm8sOBwU7lLnCpSbTyY3yeU1Vc7l4KT5zT4s/yOxHH5O7tIuuLOCnLADRw=="
},
"node_modules/d3-format": {
- "version": "3.1.0",
- "resolved": "https://registry.npmjs.org/d3-format/-/d3-format-3.1.0.tgz",
- "integrity": "sha512-YyUI6AEuY/Wpt8KWLgZHsIU86atmikuoOmCfommt0LYHiQSPjvX2AcFc38PX0CBpr2RCyZhjex+NS/LPOv6YqA==",
+ "version": "3.1.2",
+ "resolved": "https://registry.npmjs.org/d3-format/-/d3-format-3.1.2.tgz",
+ "integrity": "sha512-AJDdYOdnyRDV5b6ArilzCPPwc1ejkHcoyFarqlPqT7zRYjhavcT3uSrqcMvsgh2CgoPbK3RCwyHaVyxYcP2Arg==",
+ "license": "ISC",
"engines": {
"node": ">=12"
}
@@ -1569,6 +1580,7 @@
"version": "1.4.1",
"resolved": "https://registry.npmjs.org/dom-serializer/-/dom-serializer-1.4.1.tgz",
"integrity": "sha512-VHwB3KfrcOOkelEG2ZOfxqLZdfkil8PtJi4P8N2MMXucZq2yLp75ClViUlOVwyoHEDjYU433Aq+5zWP61+RGag==",
+ "license": "MIT",
"dependencies": {
"domelementtype": "^2.0.1",
"domhandler": "^4.2.0",
@@ -1582,6 +1594,7 @@
"version": "2.2.0",
"resolved": "https://registry.npmjs.org/entities/-/entities-2.2.0.tgz",
"integrity": "sha512-p92if5Nz619I0w+akJrLZH0MX0Pb5DX39XOwQTtXSdQQOaYH03S1uIQp4mhOZtAXrxq4ViO67YTiLBo2638o9A==",
+ "license": "BSD-2-Clause",
"funding": {
"url": "https://github.com/fb55/entities?sponsor=1"
}
@@ -1595,12 +1608,14 @@
"type": "github",
"url": "https://github.com/sponsors/fb55"
}
- ]
+ ],
+ "license": "BSD-2-Clause"
},
"node_modules/domhandler": {
"version": "4.3.1",
"resolved": "https://registry.npmjs.org/domhandler/-/domhandler-4.3.1.tgz",
"integrity": "sha512-GrwoxYN+uWlzO8uhUXRl0P+kHE4GtVPfYzVLcUxPL7KNdHKj66vvlhiweIHqYYXWlw+T8iLMp42Lm67ghw4WMQ==",
+ "license": "BSD-2-Clause",
"dependencies": {
"domelementtype": "^2.2.0"
},
@@ -1615,6 +1630,7 @@
"version": "2.8.0",
"resolved": "https://registry.npmjs.org/domutils/-/domutils-2.8.0.tgz",
"integrity": "sha512-w96Cjofp72M5IIhpjgobBimYEfoPjx1Vx0BSX9P30WBdZW2WIKU0T1Bd0kz2eNZ9ikjKgHbEyKx8BB6H1L3h3A==",
+ "license": "BSD-2-Clause",
"dependencies": {
"dom-serializer": "^1.0.1",
"domelementtype": "^2.2.0",
@@ -1693,6 +1709,7 @@
"version": "3.0.0",
"resolved": "https://registry.npmjs.org/escape-goat/-/escape-goat-3.0.0.tgz",
"integrity": "sha512-w3PwNZJwRxlp47QGzhuEBldEqVHHhh8/tIPcl6ecf2Bou99cdAt0knihBV0Ecc7CGxYduXVBDheH1K2oADRlvw==",
+ "license": "MIT",
"engines": {
"node": ">=10"
},
@@ -1704,6 +1721,7 @@
"version": "3.2.25",
"resolved": "https://registry.npmjs.org/esm/-/esm-3.2.25.tgz",
"integrity": "sha512-U1suiZ2oDVWv4zPO56S0NcR5QriEahGtdN2OR6FiOG4WJvcjBVFB0qI4+eKoWFH483PKGuLuu6V8Z4T5g63UVA==",
+ "license": "MIT",
"engines": {
"node": ">=6"
}
@@ -1799,6 +1817,7 @@
"url": "https://github.com/sponsors/fb55"
}
],
+ "license": "MIT",
"dependencies": {
"domelementtype": "^2.0.1",
"domhandler": "^4.0.0",
@@ -1810,6 +1829,7 @@
"version": "2.2.0",
"resolved": "https://registry.npmjs.org/entities/-/entities-2.2.0.tgz",
"integrity": "sha512-p92if5Nz619I0w+akJrLZH0MX0Pb5DX39XOwQTtXSdQQOaYH03S1uIQp4mhOZtAXrxq4ViO67YTiLBo2638o9A==",
+ "license": "BSD-2-Clause",
"funding": {
"url": "https://github.com/fb55/entities?sponsor=1"
}
@@ -1847,6 +1867,7 @@
"version": "8.1.0",
"resolved": "https://registry.npmjs.org/juice/-/juice-8.1.0.tgz",
"integrity": "sha512-FLzurJrx5Iv1e7CfBSZH68dC04EEvXvvVvPYB7Vx1WAuhCp1ZPIMtqxc+WTWxVkpTIC2Ach/GAv0rQbtGf6YMA==",
+ "license": "MIT",
"dependencies": {
"cheerio": "1.0.0-rc.10",
"commander": "^6.1.0",
@@ -1887,9 +1908,20 @@
}
},
"node_modules/linkify-it": {
- "version": "5.0.0",
- "resolved": "https://registry.npmjs.org/linkify-it/-/linkify-it-5.0.0.tgz",
- "integrity": "sha512-5aHCbzQRADcdP+ATqnDuhhJ/MRIqDkZX5pyjFHRRysS8vZ5AbqGEoFIb6pYHPZ+L/OC2Lc+xT8uHVVR5CAK/wQ==",
+ "version": "5.0.2",
+ "resolved": "https://registry.npmjs.org/linkify-it/-/linkify-it-5.0.2.tgz",
+ "integrity": "sha512-ONTm2jCMAVZjgQa/Fy1kScXsuOoF5NPTsoFBdE1KVIZ2vAh/r9+Bqo+0jINCBYnavTPQZz38QzFTme79ENoN3Q==",
+ "funding": [
+ {
+ "type": "github",
+ "url": "https://github.com/sponsors/puzrin"
+ },
+ {
+ "type": "github",
+ "url": "https://github.com/sponsors/markdown-it"
+ }
+ ],
+ "license": "MIT",
"dependencies": {
"uc.micro": "^2.0.0"
}
@@ -1921,13 +1953,24 @@
"integrity": "sha512-1I+1qpDt4idfgLQG+BNWmrqku+7/2bi5nLf4YwF8y8zXvmfiTBY3PV3ZibfrjBueCByROpuBjLLFCajqkgYoLQ=="
},
"node_modules/markdown-it": {
- "version": "14.1.0",
- "resolved": "https://registry.npmjs.org/markdown-it/-/markdown-it-14.1.0.tgz",
- "integrity": "sha512-a54IwgWPaeBCAAsv13YgmALOF1elABB08FxO9i+r4VFk5Vl4pKokRPeX8u5TCgSsPi6ec1otfLjdOpVcgbpshg==",
+ "version": "14.3.0",
+ "resolved": "https://registry.npmjs.org/markdown-it/-/markdown-it-14.3.0.tgz",
+ "integrity": "sha512-RCEsPjR+sr0x+AuYp601tKTkgFG4YEPLCzHST3cQ/fhlJkqAkz1L2/Qbp1j9qw5SBwQHFBoW8+hoN5xssOF0Tw==",
+ "funding": [
+ {
+ "type": "github",
+ "url": "https://github.com/sponsors/puzrin"
+ },
+ {
+ "type": "github",
+ "url": "https://github.com/sponsors/markdown-it"
+ }
+ ],
+ "license": "MIT",
"dependencies": {
"argparse": "^2.0.1",
- "entities": "^4.4.0",
- "linkify-it": "^5.0.0",
+ "entities": "^4.5.0",
+ "linkify-it": "^5.0.2",
"mdurl": "^2.0.0",
"punycode.js": "^2.3.1",
"uc.micro": "^2.1.0"
@@ -1945,15 +1988,17 @@
"version": "4.3.2",
"resolved": "https://registry.npmjs.org/markdown-it-mathjax3/-/markdown-it-mathjax3-4.3.2.tgz",
"integrity": "sha512-TX3GW5NjmupgFtMJGRauioMbbkGsOXAAt1DZ/rzzYmTHqzkO1rNAdiMD4NiruurToPApn2kYy76x02QN26qr2w==",
+ "license": "MIT",
"dependencies": {
"juice": "^8.0.0",
"mathjax-full": "^3.2.0"
}
},
"node_modules/mathjax-full": {
- "version": "3.2.2",
- "resolved": "https://registry.npmjs.org/mathjax-full/-/mathjax-full-3.2.2.tgz",
- "integrity": "sha512-+LfG9Fik+OuI8SLwsiR02IVdjcnRCy5MufYLi0C3TdMT56L/pjB0alMVGgoWJF8pN9Rc7FESycZB9BMNWIid5w==",
+ "version": "3.2.1",
+ "resolved": "https://registry.npmjs.org/mathjax-full/-/mathjax-full-3.2.1.tgz",
+ "integrity": "sha512-aUz9o16MGZdeiIBwZjAfUBTiJb7LRqzZEl1YOZ8zQMGYIyh1/nxRebxKxjDe9L+xcZCr2OHdzoFBMcd6VnLv9Q==",
+ "license": "Apache-2.0",
"dependencies": {
"esm": "^3.2.25",
"mhchemparser": "^4.1.0",
@@ -1989,12 +2034,14 @@
"node_modules/mensch": {
"version": "0.3.4",
"resolved": "https://registry.npmjs.org/mensch/-/mensch-0.3.4.tgz",
- "integrity": "sha512-IAeFvcOnV9V0Yk+bFhYR07O3yNina9ANIN5MoXBKYJ/RLYPurd2d0yw14MDhpr9/momp0WofT1bPUh3hkzdi/g=="
+ "integrity": "sha512-IAeFvcOnV9V0Yk+bFhYR07O3yNina9ANIN5MoXBKYJ/RLYPurd2d0yw14MDhpr9/momp0WofT1bPUh3hkzdi/g==",
+ "license": "MIT"
},
"node_modules/mhchemparser": {
"version": "4.2.1",
"resolved": "https://registry.npmjs.org/mhchemparser/-/mhchemparser-4.2.1.tgz",
- "integrity": "sha512-kYmyrCirqJf3zZ9t/0wGgRZ4/ZJw//VwaRVGA75C4nhE60vtnIzhl9J9ndkX/h6hxSN7pjg/cE0VxbnNM+bnDQ=="
+ "integrity": "sha512-kYmyrCirqJf3zZ9t/0wGgRZ4/ZJw//VwaRVGA75C4nhE60vtnIzhl9J9ndkX/h6hxSN7pjg/cE0VxbnNM+bnDQ==",
+ "license": "Apache-2.0"
},
"node_modules/micromark-util-character": {
"version": "2.1.1",
@@ -2105,7 +2152,8 @@
"node_modules/mj-context-menu": {
"version": "0.6.1",
"resolved": "https://registry.npmjs.org/mj-context-menu/-/mj-context-menu-0.6.1.tgz",
- "integrity": "sha512-7NO5s6n10TIV96d4g2uDpG7ZDpIhMh0QNfGdJw/W47JswFcosz457wqz/b5sAKvl12sxINGFCn80NZHKwxQEXA=="
+ "integrity": "sha512-7NO5s6n10TIV96d4g2uDpG7ZDpIhMh0QNfGdJw/W47JswFcosz457wqz/b5sAKvl12sxINGFCn80NZHKwxQEXA==",
+ "license": "Apache-2.0"
},
"node_modules/nanoid": {
"version": "3.3.11",
@@ -2144,6 +2192,7 @@
"version": "2.7.0",
"resolved": "https://registry.npmjs.org/node-fetch/-/node-fetch-2.7.0.tgz",
"integrity": "sha512-c4FRfUm/dbcWZ7U+1Wq0AwCyFL+3nt2bEw05wfxSz+DWpWsitgmSgYmy2dQdWyKC1694ELPqMs/YzUSNozLt8A==",
+ "license": "MIT",
"dependencies": {
"whatwg-url": "^5.0.0"
},
@@ -2163,6 +2212,7 @@
"version": "2.1.1",
"resolved": "https://registry.npmjs.org/nth-check/-/nth-check-2.1.1.tgz",
"integrity": "sha512-lqjrjmaOoAnWfMmBPL+XNnynZh2+swxiX3WUE0s4yEHI6m+AwrK2UZOimIRl3X/4QctVqS8AiZjFqyOGrMXb/w==",
+ "license": "BSD-2-Clause",
"dependencies": {
"boolbase": "^1.0.0"
},
@@ -2191,12 +2241,14 @@
"node_modules/parse5": {
"version": "6.0.1",
"resolved": "https://registry.npmjs.org/parse5/-/parse5-6.0.1.tgz",
- "integrity": "sha512-Ofn/CTFzRGTTxwpNEs9PP93gXShHcTq255nzRYSKe8AkVpZY7e1fpmTfOyoIvjP5HG7Z2ZM7VS9PPhQGW2pOpw=="
+ "integrity": "sha512-Ofn/CTFzRGTTxwpNEs9PP93gXShHcTq255nzRYSKe8AkVpZY7e1fpmTfOyoIvjP5HG7Z2ZM7VS9PPhQGW2pOpw==",
+ "license": "MIT"
},
"node_modules/parse5-htmlparser2-tree-adapter": {
"version": "6.0.1",
"resolved": "https://registry.npmjs.org/parse5-htmlparser2-tree-adapter/-/parse5-htmlparser2-tree-adapter-6.0.1.tgz",
"integrity": "sha512-qPuWvbLgvDGilKc5BoicRovlT4MtYT6JfJyBOMDsKoiT+GiuP5qyrPCnR9HcPECIJJmZh5jRndyNThnhhb/vlA==",
+ "license": "MIT",
"dependencies": {
"parse5": "^6.0.1"
}
@@ -2391,6 +2443,7 @@
"version": "1.12.2",
"resolved": "https://registry.npmjs.org/slick/-/slick-1.12.2.tgz",
"integrity": "sha512-4qdtOGcBjral6YIBCWJ0ljFSKNLz9KkhbWtuGvUyRowl1kxfuE1x/Z/aJcaiilpb3do9bl5K7/1h9XC5wWpY/A==",
+ "license": "MIT (http://mootools.net/license.txt)",
"engines": {
"node": "*"
}
@@ -2430,11 +2483,12 @@
}
},
"node_modules/speech-rule-engine": {
- "version": "4.1.2",
- "resolved": "https://registry.npmjs.org/speech-rule-engine/-/speech-rule-engine-4.1.2.tgz",
- "integrity": "sha512-S6ji+flMEga+1QU79NDbwZ8Ivf0S/MpupQQiIC0rTpU/ZTKgcajijJJb1OcByBQDjrXCN1/DJtGz4ZJeBMPGJw==",
+ "version": "4.1.4",
+ "resolved": "https://registry.npmjs.org/speech-rule-engine/-/speech-rule-engine-4.1.4.tgz",
+ "integrity": "sha512-i/VCLG1fvRc95pMHRqG4aQNscv+9aIsqA2oI7ZQS51sTdUcDHYX6cpT8/tqZ+enjs1tKVwbRBWgxut9SWn+f9g==",
+ "license": "Apache-2.0",
"dependencies": {
- "@xmldom/xmldom": "0.9.8",
+ "@xmldom/xmldom": "0.9.10",
"commander": "13.1.0",
"wicked-good-xpath": "1.3.0"
},
@@ -2446,6 +2500,7 @@
"version": "13.1.0",
"resolved": "https://registry.npmjs.org/commander/-/commander-13.1.0.tgz",
"integrity": "sha512-/rFeCpNJQbhSZjGVwO9RFV3xPqbnERS8MmIQzCtD/zl6gpJuV/bMLuN92oG3F7d8oDEHHRrujSXNUr8fpjntKw==",
+ "license": "MIT",
"engines": {
"node": ">=18"
}
@@ -2507,7 +2562,8 @@
"node_modules/tr46": {
"version": "0.0.3",
"resolved": "https://registry.npmjs.org/tr46/-/tr46-0.0.3.tgz",
- "integrity": "sha512-N3WMsuqV66lT30CrXNbEjx4GEwlow3v6rr4mCcv6prnfwhS01rkgyFdjPNBYd9br7LpXV1+Emh01fHnq2Gdgrw=="
+ "integrity": "sha512-N3WMsuqV66lT30CrXNbEjx4GEwlow3v6rr4mCcv6prnfwhS01rkgyFdjPNBYd9br7LpXV1+Emh01fHnq2Gdgrw==",
+ "license": "MIT"
},
"node_modules/trim-lines": {
"version": "3.0.1",
@@ -2526,7 +2582,8 @@
"node_modules/uc.micro": {
"version": "2.1.0",
"resolved": "https://registry.npmjs.org/uc.micro/-/uc.micro-2.1.0.tgz",
- "integrity": "sha512-ARDJmphmdvUk6Glw7y9DQ2bFkKBHwQHLi2lsaH6PPmz/Ka9sFOBsBluozhDltWmnv9u/cF6Rt87znRTPV+yp/A=="
+ "integrity": "sha512-ARDJmphmdvUk6Glw7y9DQ2bFkKBHwQHLi2lsaH6PPmz/Ka9sFOBsBluozhDltWmnv9u/cF6Rt87znRTPV+yp/A==",
+ "license": "MIT"
},
"node_modules/unist-util-is": {
"version": "6.0.0",
@@ -2595,6 +2652,7 @@
"version": "3.0.1",
"resolved": "https://registry.npmjs.org/valid-data-url/-/valid-data-url-3.0.1.tgz",
"integrity": "sha512-jOWVmzVceKlVVdwjNSenT4PbGghU0SBIizAev8ofZVgivk/TVHXSbNL8LP6M3spZvkR9/QolkyJavGSX5Cs0UA==",
+ "license": "MIT",
"engines": {
"node": ">=10"
}
@@ -2724,12 +2782,17 @@
}
},
"node_modules/vitepress-plugin-tabs": {
- "version": "0.6.0",
- "resolved": "https://registry.npmjs.org/vitepress-plugin-tabs/-/vitepress-plugin-tabs-0.6.0.tgz",
- "integrity": "sha512-35mQUbI6lcVrLj0PpwautXm8LiV6Gz5ZDlklVRf/0KDd6gGs8HaZztqOxdffnzCiKIFLURBx930xq28it9jQTg==",
+ "version": "0.9.0",
+ "resolved": "https://registry.npmjs.org/vitepress-plugin-tabs/-/vitepress-plugin-tabs-0.9.0.tgz",
+ "integrity": "sha512-OB+/T6SkwGIQHIQi1NLl011kWCMzfsVEx0/q+NfF6EmvvpcfDzRClc/iEwUlq8TkQnFxANu+GTdmvtXEjJM2Lg==",
+ "license": "MIT",
+ "funding": {
+ "type": "github",
+ "url": "https://github.com/sponsors/sapphi-red"
+ },
"peerDependencies": {
- "vitepress": "^1.0.0",
- "vue": "^3.3.8"
+ "vitepress": "^1.0.0 || ^2.0.0-alpha.17",
+ "vue": "^3.5.0"
}
},
"node_modules/vue": {
@@ -2756,6 +2819,7 @@
"version": "6.0.1",
"resolved": "https://registry.npmjs.org/web-resource-inliner/-/web-resource-inliner-6.0.1.tgz",
"integrity": "sha512-kfqDxt5dTB1JhqsCUQVFDj0rmY+4HLwGQIsLPbyrsN9y9WV/1oFDSx3BQ4GfCv9X+jVeQ7rouTqwK53rA/7t8A==",
+ "license": "MIT",
"dependencies": {
"ansi-colors": "^4.1.1",
"escape-goat": "^3.0.0",
@@ -2772,6 +2836,7 @@
"version": "3.3.0",
"resolved": "https://registry.npmjs.org/domhandler/-/domhandler-3.3.0.tgz",
"integrity": "sha512-J1C5rIANUbuYK+FuFL98650rihynUOEzRLxW+90bKZRWB6A1X1Tf82GxR1qAWLyfNPRvjqfip3Q5tdYlmAa9lA==",
+ "license": "BSD-2-Clause",
"dependencies": {
"domelementtype": "^2.0.1"
},
@@ -2786,6 +2851,7 @@
"version": "2.2.0",
"resolved": "https://registry.npmjs.org/entities/-/entities-2.2.0.tgz",
"integrity": "sha512-p92if5Nz619I0w+akJrLZH0MX0Pb5DX39XOwQTtXSdQQOaYH03S1uIQp4mhOZtAXrxq4ViO67YTiLBo2638o9A==",
+ "license": "BSD-2-Clause",
"funding": {
"url": "https://github.com/fb55/entities?sponsor=1"
}
@@ -2794,6 +2860,7 @@
"version": "5.0.1",
"resolved": "https://registry.npmjs.org/htmlparser2/-/htmlparser2-5.0.1.tgz",
"integrity": "sha512-vKZZra6CSe9qsJzh0BjBGXo8dvzNsq/oGvsjfRdOrrryfeD9UOBEEQdeoqCRmKZchF5h2zOBMQ6YuQ0uRUmdbQ==",
+ "license": "MIT",
"dependencies": {
"domelementtype": "^2.0.1",
"domhandler": "^3.3.0",
@@ -2808,6 +2875,7 @@
"version": "2.6.0",
"resolved": "https://registry.npmjs.org/mime/-/mime-2.6.0.tgz",
"integrity": "sha512-USPkMeET31rOMiarsBNIHZKLGgvKc/LrjofAnBlOttf5ajRvqiRA8QsenbcooctK6d6Ts6aqZXBA+XbkKthiQg==",
+ "license": "MIT",
"bin": {
"mime": "cli.js"
},
@@ -2818,12 +2886,14 @@
"node_modules/webidl-conversions": {
"version": "3.0.1",
"resolved": "https://registry.npmjs.org/webidl-conversions/-/webidl-conversions-3.0.1.tgz",
- "integrity": "sha512-2JAn3z8AR6rjK8Sm8orRC0h/bcl/DqL7tRPdGZ4I1CjdF+EaMLmYxBHyXuKL849eucPFhvBoxMsflfOb8kxaeQ=="
+ "integrity": "sha512-2JAn3z8AR6rjK8Sm8orRC0h/bcl/DqL7tRPdGZ4I1CjdF+EaMLmYxBHyXuKL849eucPFhvBoxMsflfOb8kxaeQ==",
+ "license": "BSD-2-Clause"
},
"node_modules/whatwg-url": {
"version": "5.0.0",
"resolved": "https://registry.npmjs.org/whatwg-url/-/whatwg-url-5.0.0.tgz",
"integrity": "sha512-saE57nupxk6v3HY35+jzBwYa0rKSy0XR8JSxZPwgLr7ys0IBzhGviA1/TUGJLmSVqs8pb9AnvICXEuOHLprYTw==",
+ "license": "MIT",
"dependencies": {
"tr46": "~0.0.3",
"webidl-conversions": "^3.0.0"
@@ -2832,7 +2902,8 @@
"node_modules/wicked-good-xpath": {
"version": "1.3.0",
"resolved": "https://registry.npmjs.org/wicked-good-xpath/-/wicked-good-xpath-1.3.0.tgz",
- "integrity": "sha512-Gd9+TUn5nXdwj/hFsPVx5cuHHiF5Bwuc30jZ4+ronF1qHK5O7HD0sgmXWSEgwKquT3ClLoKPVbO6qGwVwLzvAw=="
+ "integrity": "sha512-Gd9+TUn5nXdwj/hFsPVx5cuHHiF5Bwuc30jZ4+ronF1qHK5O7HD0sgmXWSEgwKquT3ClLoKPVbO6qGwVwLzvAw==",
+ "license": "MIT"
},
"node_modules/zwitch": {
"version": "2.0.4",
diff --git a/docs/package.json b/docs/package.json
index dfec562ea..e92aaa1df 100644
--- a/docs/package.json
+++ b/docs/package.json
@@ -11,11 +11,11 @@
},
"dependencies": {
"@nolebase/vitepress-plugin-enhanced-readabilities": "^2.14.0",
- "d3-format": "^3.1.0",
- "markdown-it": "^14.1.0",
+ "d3-format": "^3.1.2",
+ "markdown-it": "^14.2.0",
"markdown-it-footnote": "^4.0.0",
"markdown-it-mathjax3": "^4.3.2",
"vitepress": "^1.6.3",
- "vitepress-plugin-tabs": "^0.6.0"
+ "vitepress-plugin-tabs": "^0.9.0"
}
}
diff --git a/docs/papers.yml b/docs/papers.yml
index 21afd44fd..2709f1152 100644
--- a/docs/papers.yml
+++ b/docs/papers.yml
@@ -2,6 +2,35 @@
# information to generate the "Research Showcase"
papers:
+ - title: "Maximum updraft velocity beyond CAPE: the role of boundary layer dynamics and pressure perturbations"
+ authors:
+ - Bety Pechacova (1)
+ - Alejandro Casallas (1)
+ - Tom Beucler (2,3)
+ - Lokahith Agasthya (4)
+ - Caroline Muller (1)
+ affiliations:
+ 1: Institute of Science and Technology Austria
+ 2: Faculty of Geosciences and Environment, University of Lausanne
+ 3: Expertise Center for Climate Extremes, University of Lausanne
+ 4: Department of Meteorology and Geophysics, University of Vienna
+ link: https://doi.org/10.48550/arXiv.2607.23627
+ abstract: "Deep convective updraft velocities play a key role in the Earth's climate system, influencing precipitation extremes, lightning, and the planetary energy budget. While Convective Available Potential Energy (CAPE) is widely used to explain maximum updraft velocity ($w_{\\max}$), CAPE is an imperfect predictor as updrafts are also influenced by entrainment, boundary layer dynamics, pressure perturbations, and condensate loading. However, the relative importance of these processes and how they interact to set $w_{\\max}$ in individual clouds remains unclear. Here, we use equation learning to identify compact, physically interpretable relationships linking environmental and in-cloud conditions to $w_{\\max}$ in individual tracked clouds across idealized radiative–convective equilibrium regimes spanning a range of sea surface temperatures and radiative cooling rates. For pre-storm prediction, CAPE and local mean boundary layer vertical velocity ($\\overline{w_{\\mathrm{bl}}}$) together explain nearly half the variance in $w_{\\max}$ across regimes ($R^2=0.47$). While CAPE captures regime-mean differences, it has little predictive value within a single simulation. $\\overline{w_{\\mathrm{bl}}}$ is essential for capturing cloud-to-cloud variability, including the suppression of $w_{\\max}$ even at high $\\mathrm{CAPE}$ values. At the time of peak intensity, a simple approximate Bernoulli-like invariant combining maximum pressure perturbation and maximum cloud condensate explains 89% of the variance ($R^2=0.89$). The tight link between $w_{\\max}$ and pressure perturbation supports the sticky thermals hypothesis and highlights the importance of dynamic pressure effects, often neglected in updraft theories. These results highlight $\\overline{w_{\\mathrm{bl}}}$ as an important regulator of convective intensity alongside CAPE, and demonstrate that dynamic pressure plays an important role within individual updrafts."
+ image: https://raw.githubusercontent.com/MilesCranmer/PySR_Docs/d6a8df0a069efb549ed9fc53e87cf0877ad3fee6/images/pr1295_updraft_velocity.png
+ date: 2026-07-26
+ - title: Improvement of spatiotemporal generalization in radiocesium transfer models for wheat using symbolic regression
+ authors:
+ - Katsutoshi Seki (1)
+ - Noriko Yamaguchi (2)
+ - Tetsuya Eguchi (2)
+ - Masato Igura (2)
+ affiliations:
+ 1: Natural Science Laboratory, Toyo University
+ 2: Institute for Agro-Environmental Sciences, NARO
+ link: https://doi.org/10.1016/j.jenvrad.2026.108077
+ abstract: "Accurate prediction of the soil-to-plant transfer factor (TF) of radiocesium ($^{137}$Cs) is essential for assessing radionuclide transfer to the human food chain and supporting radiological risk management after nuclear contamination. Semi-empirical TF models based on exchangeable potassium (K) and soil fixation capacity are widely used, but their predictive performance often deteriorates when applied to unobserved sites or years. In this study, we evaluated whether symbolic regression (SR), an approach that derives interpretable mathematical equations from data, can improve the spatiotemporal generalization of TF models. We analyzed long-term monitoring data from 11 upland wheat fields in Japan collected over five survey years ($n$ = 36). Model performance was assessed using structured cross-validation, including leave-one-site-out (LOSO) and leave-one-year-out (LOYO), to explicitly test spatial and temporal extrapolation. A simple K-based model substantially improved baseline predictions compared with a conventional semi-empirical model. However, linear additive extensions incorporating radiocesium interception potential (RIP), cation exchange capacity, pH, and soil $^{137}$Cs did not consistently improve predictive performance under LOSO or LOYO. In contrast, the selected SR model showed the lowest point-estimate prediction errors in both spatial and temporal validations. The selected equation represented a non-linear interaction between exchangeable K and RIP, consistent with reduced K sensitivity in soils with strong specific fixation sites. However, the selected SR-derived equation remains empirical and is valid only within the observed soil-property range and may show nonphysical behavior under extreme RIP conditions. These results suggest that interpretable non-linear equations derived by SR may improve TF prediction within a restricted soil domain."
+ image: https://raw.githubusercontent.com/MilesCranmer/PySR_Docs/d6a8df0a069efb549ed9fc53e87cf0877ad3fee6/images/pr1206_radiocesium_transfer_models.png
+ date: 2026-06-13
- title: Discovering data-driven microbial growth models with symbolic regression
authors:
- T. Anthony Sun (1)
@@ -80,7 +109,7 @@ papers:
- Pericles Panagiotou (1)
affiliations:
1: Laboratory of Fluid Mechanics and Turbomachinery, Department of Mechanical Engineering, Aristotle University of Thessaloniki, 54124 Thessaloniki, Greece
- link: https://www.aifluids.net/proceedings/S6P16.pdf
+ link: https://www.aifluids.net/proceedings/C.%20Pliakos,%20G.%20Efrem,%20D.%20Terzis,%20P.%20Panagiotou,%20Data-driven%20skin%20friction%20estimation%20for%20UAV%20wings%20in%20subsonic%20flows.pdf
abstract: Accurate estimation of the skin friction coefficient (𝐶𝑓) is essential for estimating the wall shear stresses (𝜏𝑤) and ultimately the first-layer cell height (𝑦) in wall-resolved RANS simulations of wings, where turbulence models are used, demanding a specific grid resolution near walls (primarily the 𝑦𝑡𝑎𝑟𝑔𝑒𝑡⁺). Conventional flat-plate correlations often fail to account for the three-dimensional nature of real wing flows, introducing uncertainties in 𝐶𝑓 predictions and leading to multiple CFD analyses and mesh refinements to meet the targets. In this work, we propose a machine-learning-based approach exploring symbolic regression to derive a model that correlates wing-specific parameters (e.g., Reynolds number, angle of attack, thickness-to-chord ratio, wing sweep angle) with 𝐶𝑓 at the Mean Aerodynamic Chord (MAC). Data are acquired from an in-house database of over 5,000 RANS simulations for UAV wings operating in the low subsonic regime, covering a wide design space, all conducted following best-practice CFD guidelines to ensure high fidelity. These analyses are performed at various flow conditions covering Reynolds numbers from 10⁵ to 10⁷ and include the complete drag polar for each wing. The proposed correlation provides improved agreement with CFD data and enables more accurate 𝑦⁺ estimations. Validation on different wing geometries, including the ONERA M6 and in-house UAV wings, confirmed the robustness of the model, which improves boundary-layer resolution with only a marginal (~2%) increase in total mesh size, while achieving an R² of 0.68 with negligible computational inference cost. This explicit, data-driven equation offers an efficient method for streamlining mesh generation in aerodynamic simulations.
image: https://raw.githubusercontent.com/MilesCranmer/PySR_Docs/48020babcca1e73d2f5bbbd4ef3f0ec3265c59fb/images/pr1127_skin_friction.png
date: 2025-05-30
diff --git a/docs/src/.vitepress/config.mts b/docs/src/.vitepress/config.mts
index 1921fdd00..7572d3fda 100644
--- a/docs/src/.vitepress/config.mts
+++ b/docs/src/.vitepress/config.mts
@@ -1,20 +1,31 @@
+import { readFileSync } from 'node:fs'
+import path from 'node:path'
+import { fileURLToPath } from 'node:url'
import { defineConfig } from 'vitepress'
import { tabsMarkdownPlugin } from 'vitepress-plugin-tabs'
import mathjax3 from "markdown-it-mathjax3";
import footnote from "markdown-it-footnote";
-function getBaseRepository(base: string): string {
- if (!base || base === '/') return '/';
- const parts = base.split('/').filter(Boolean);
- return parts.length > 0 ? `/${parts[0]}/` : '/';
+// Maps the juliapkg.json pin to the SymbolicRegression.jl docs folder covering
+// every version the pin admits, following Julia Pkg compat semantics:
+// `=X.Y.Z` is exact, `~X.Y.Z` allows patch bumps, `^X.Y.Z` (or bare) allows minor bumps.
+export function symbolicRegressionDocsSubfolder(pkg: { version?: unknown; rev?: unknown }): string {
+ const compat = typeof pkg.version === 'string' ? pkg.version.trim().match(/^([=~^]?)(\d+)(?:\.(\d+))?(?:\.(\d+))?$/) : null
+ if (compat) {
+ const [, op, major, minor, patch] = compat
+ if (op === '=' && patch !== undefined) return `v${major}.${minor}.${patch}`
+ if ((op === '~' || op === '=' || major === '0') && minor !== undefined) return `v${major}.${minor}`
+ return `v${major}`
+ }
+ return typeof pkg.rev === 'string' && /^v\d+\.\d+\.\d+$/.test(pkg.rev) ? pkg.rev : 'dev'
}
-// Canonical domain for SEO - will be replaced by deploy.jl at build time
-// Empty string means no canonical tag (this site is canonical)
-const canonicalDomain = '';
+const configDirectory = path.dirname(fileURLToPath(import.meta.url))
+const juliaPkg = JSON.parse(readFileSync(path.resolve(configDirectory, '../../../pysr/juliapkg.json'), 'utf8'))
+const symbolicRegressionDocsVersion = symbolicRegressionDocsSubfolder(juliaPkg.packages.SymbolicRegression)
const baseTemp = {
- base: '/pysr/',
+ base: '/',
}
const nav = [
@@ -23,13 +34,13 @@ const nav = [
{ text: 'API', link: '/api' },
{
text: 'Forum',
- link: 'https://github.com/MilesCranmer/PySR/discussions'
+ link: 'https://github.com/astroautomata/PySR/discussions'
},
{
text: 'Python',
items: [
{ text: 'Python', link: '/' },
- { text: 'Julia', link: 'https://ai.damtp.cam.ac.uk/symbolicregression/dev/' }
+ { text: 'Julia', link: `https://julia.pysr.ai/${symbolicRegressionDocsVersion}/`, target: '_self', noIcon: true }
]
},
{
@@ -39,7 +50,7 @@ const nav = [
// https://vitepress.dev/reference/site-config
export default defineConfig({
- base: '/pysr/',
+ base: '/',
title: 'PySR',
description: 'High-Performance Symbolic Regression in Python and Julia',
lastUpdated: true,
@@ -47,25 +58,17 @@ export default defineConfig({
outDir: '../dist',
srcExclude: ['**/_*.md'],
head: [
- ['link', { rel: 'icon', href: `${baseTemp.base}favicon.png` }],
- ['script', {src: `${getBaseRepository(baseTemp.base)}versions.js`}],
+ ['link', { rel: 'icon', type: 'image/x-icon', href: `${baseTemp.base}favicon.ico` }],
+ ['link', { rel: 'icon', type: 'image/png', sizes: '16x16', href: `${baseTemp.base}favicon-16x16.png` }],
+ ['link', { rel: 'icon', type: 'image/png', sizes: '32x32', href: `${baseTemp.base}favicon-32x32.png` }],
+ ['link', { rel: 'apple-touch-icon', sizes: '180x180', href: `${baseTemp.base}apple-touch-icon.png` }],
+ ['script', {src: '/versions.js'}],
['script', {src: `${baseTemp.base}siteinfo.js`}]
],
ignoreDeadLinks: true,
- transformHead: ({ pageData }) => {
- if (!canonicalDomain) return [];
-
- // Construct canonical URL for this page
- const pagePath = pageData.relativePath.replace(/((^|\/)index)?\.md$/, '$2');
- const canonicalUrl = `${canonicalDomain}${pagePath}`;
-
- return [
- ['link', { rel: 'canonical', href: canonicalUrl }]
- ];
- },
vite: {
define: {
- __DEPLOY_ABSPATH__: JSON.stringify(getBaseRepository(baseTemp.base)),
+ __DEPLOY_ABSPATH__: JSON.stringify('/'),
},
optimizeDeps: {
exclude: [
@@ -96,7 +99,7 @@ export default defineConfig({
},
themeConfig: {
outline: 'deep',
- logo: 'https://ai.damtp.cam.ac.uk/symbolicregression/dev/logo.png',
+ logo: 'https://raw.githubusercontent.com/MilesCranmer/PySR_Docs/master/images/pysr-logo.png',
search: {
provider: 'local',
options: {
@@ -109,7 +112,21 @@ export default defineConfig({
text: 'Getting Started',
items: [
{ text: 'Introduction', link: '/' },
- { text: 'Examples', link: '/examples' },
+ {
+ text: 'Examples',
+ link: '/examples',
+ collapsed: true,
+ items: [
+ { text: 'Getting started', link: '/examples/getting-started' },
+ { text: 'Expression specifications', link: '/examples/expression-specifications' },
+ { text: 'Objectives and losses', link: '/examples/objectives' },
+ { text: 'Physics and units', link: '/examples/physics' },
+ { text: 'Search behaviour', link: '/examples/search-behaviour' },
+ { text: 'Instrumentation and workflow', link: '/examples/instrumentation' },
+ { text: 'Value types', link: '/examples/value-types' },
+ { text: 'Beyond numeric values', link: '/examples/beyond-numeric-values' },
+ ]
+ },
]
},
{
@@ -131,16 +148,18 @@ export default defineConfig({
text: 'Advanced',
items: [
{ text: 'Tuning', link: '/tuning' },
+ { text: 'Slurm', link: '/slurm' },
{ text: 'Backend', link: '/backend' },
+ { text: 'Migrating from v1', link: '/migration' },
]
}
],
editLink: {
- pattern: 'https://github.com/MilesCranmer/PySR/edit/master/docs/:path',
+ pattern: 'https://github.com/astroautomata/PySR/edit/master/docs/:path',
text: 'Edit this page on GitHub'
},
socialLinks: [
- { icon: 'github', link: 'https://github.com/MilesCranmer/PySR' }
+ { icon: 'github', link: 'https://github.com/astroautomata/PySR' }
],
footer: {
message: 'Made with VitePress',
diff --git a/docs/src/.vitepress/theme/style.css b/docs/src/.vitepress/theme/style.css
index e5b79d3d1..d9afca052 100644
--- a/docs/src/.vitepress/theme/style.css
+++ b/docs/src/.vitepress/theme/style.css
@@ -108,3 +108,14 @@ figure img {
.VPNavBarMenuGroup:hover button .text {
color: var(--vp-c-text-2) !important;
}
+
+/* A \s*(.*?)\s*
\s*(.*?)\s* section whose holds a heading: keep them on one line */
+.vp-doc summary:has(> h3) {
+ cursor: pointer;
+ margin-top: 32px;
+}
+
+.vp-doc summary > h3 {
+ display: inline;
+ margin: 0;
+}
diff --git a/docs/src/_api_template.md b/docs/src/_api_template.md
index 28cfa3155..2e25d7e86 100644
--- a/docs/src/_api_template.md
+++ b/docs/src/_api_template.md
@@ -75,13 +75,15 @@ PARAMSKEY
heading_level: 3
show_root_full_path: false
-::: pysr.ParametricExpressionSpec
+::: pysr.AbstractExpressionSpec
options:
show_root_heading: true
heading_level: 3
show_root_full_path: false
-::: pysr.AbstractExpressionSpec
+## Type Specifications
+
+::: pysr.TypeSpec
options:
show_root_heading: true
heading_level: 3
@@ -100,3 +102,121 @@ PARAMSKEY
show_root_heading: true
heading_level: 3
show_root_full_path: false
+
+## Mutations
+
+::: pysr.ConstantMutation
+ options:
+ show_root_heading: true
+ heading_level: 3
+ show_root_full_path: false
+
+::: pysr.OperatorMutation
+ options:
+ show_root_heading: true
+ heading_level: 3
+ show_root_full_path: false
+
+::: pysr.FeatureMutation
+ options:
+ show_root_heading: true
+ heading_level: 3
+ show_root_full_path: false
+
+::: pysr.SwapOperandsMutation
+ options:
+ show_root_heading: true
+ heading_level: 3
+ show_root_full_path: false
+
+::: pysr.AddNodeMutation
+ options:
+ show_root_heading: true
+ heading_level: 3
+ show_root_full_path: false
+
+::: pysr.InsertNodeMutation
+ options:
+ show_root_heading: true
+ heading_level: 3
+ show_root_full_path: false
+
+::: pysr.DeleteNodeMutation
+ options:
+ show_root_heading: true
+ heading_level: 3
+ show_root_full_path: false
+
+::: pysr.RotateTreeMutation
+ options:
+ show_root_heading: true
+ heading_level: 3
+ show_root_full_path: false
+
+::: pysr.BacksolveMutation
+ options:
+ show_root_heading: true
+ heading_level: 3
+ show_root_full_path: false
+
+::: pysr.SimplifyMutation
+ options:
+ show_root_heading: true
+ heading_level: 3
+ show_root_full_path: false
+
+::: pysr.RandomizeMutation
+ options:
+ show_root_heading: true
+ heading_level: 3
+ show_root_full_path: false
+
+::: pysr.OptimizeMutation
+ options:
+ show_root_heading: true
+ heading_level: 3
+ show_root_full_path: false
+
+::: pysr.DoNothingMutation
+ options:
+ show_root_heading: true
+ heading_level: 3
+ show_root_full_path: false
+
+::: pysr.AbstractMutation
+ options:
+ show_root_heading: true
+ heading_level: 3
+ show_root_full_path: false
+
+## Plugins
+
+::: pysr.AdaptiveMutationWeightsPlugin
+ options:
+ show_root_heading: true
+ heading_level: 3
+ show_root_full_path: false
+
+::: pysr.AdaptiveParsimonyPlugin
+ options:
+ show_root_heading: true
+ heading_level: 3
+ show_root_full_path: false
+
+::: pysr.MutationBurstPlugin
+ options:
+ show_root_heading: true
+ heading_level: 3
+ show_root_full_path: false
+
+::: pysr.SimulatedAnnealingPlugin
+ options:
+ show_root_heading: true
+ heading_level: 3
+ show_root_full_path: false
+
+::: pysr.AbstractPlugin
+ options:
+ show_root_heading: true
+ heading_level: 3
+ show_root_full_path: false
diff --git a/docs/src/_index.md b/docs/src/_index.md
index 52587ba3b..df6abef05 100644
--- a/docs/src/_index.md
+++ b/docs/src/_index.md
@@ -16,7 +16,7 @@ hero:
text: View on GitHub
link: https://github.com/astroautomata/PySR
image:
- src: https://ai.damtp.cam.ac.uk/symbolicregression/dev/logo.png
+ src: https://raw.githubusercontent.com/MilesCranmer/PySR_Docs/master/images/pysr-logo.png
alt: PySR
features:
@@ -33,8 +33,8 @@ features:
details: "Customize everything: operators, loss functions, complexity, input types, optimizer, and more."
- icon: 🐍
- title: Python Native
- details: Built for Python with Julia backend for maximum performance and ease of use.
+ title: Familiar Python API
+ details: A scikit-learn interface, with results exported to SymPy, PyTorch, or JAX.
---
If you find PySR useful, please cite the paper [arXiv:2305.01582](https://arxiv.org/abs/2305.01582).
diff --git a/docs/src/backend.md b/docs/src/backend.md
index 801f353cb..3dc4b68cb 100644
--- a/docs/src/backend.md
+++ b/docs/src/backend.md
@@ -33,9 +33,9 @@ The main search code can be found in `src/SymbolicRegression.jl`.
Here are some tips:
-- The documentation for the backend is given [here](https://ai.damtp.cam.ac.uk/symbolicregression/dev/).
+- The documentation for the backend is given [here](https://julia.pysr.ai/dev/).
- Throughout the package, you will often see template functions which typically use a symbol `T` (such as in the string `where {T<:Real}`). Here, `T` is simply the datatype of the input data and stored constants, such as `Float32` or `Float64`. Writing functions in this way lets us write functions generic to types, while still having access to the specific type specified at compilation time.
-- Expressions are stored as binary trees, using the `Node{T}` type, described [here](https://ai.damtp.cam.ac.uk/symbolicregression/dev/types/#SymbolicRegression.CoreModule.EquationModule.Node).
+- Expressions are stored as binary trees, using the `Node{T}` type, described [here](https://julia.pysr.ai/dev/types#SymbolicRegression.CoreModule.EquationModule.Node).
- For reference, the main loop itself is found in the `equation_search` function inside [`src/SymbolicRegression.jl`](https://github.com/astroautomata/SymbolicRegression.jl/blob/master/src/SymbolicRegression.jl).
- Parts of the code which are typically edited by users include:
- [`src/CheckConstraints.jl`](https://github.com/astroautomata/SymbolicRegression.jl/blob/master/src/CheckConstraints.jl), particularly the function `check_constraints`. This function checks whether a given expression satisfies constraints, such as having a complexity lower than `maxsize`, and whether it contains any forbidden nestings of functions.
@@ -70,6 +70,6 @@ For more information on `juliapkg.json`, see [`pyjuliapkg`](https://github.com/J
## Additional notes
-If you get comfortable enough with the backend, you might consider using the Julia package directly: the API is given on the [SymbolicRegression.jl documentation](https://ai.damtp.cam.ac.uk/symbolicregression/dev/).
+If you get comfortable enough with the backend, you might consider using the Julia package directly: the API is given on the [SymbolicRegression.jl documentation](https://julia.pysr.ai/dev/).
If you make a change that you think could be useful to other users, don't hesitate to open a pull request on either the PySR or SymbolicRegression.jl repositories! Contributions are very appreciated.
diff --git a/docs/src/components/VersionPicker.vue b/docs/src/components/VersionPicker.vue
index 17514e206..2dc8a6b49 100644
--- a/docs/src/components/VersionPicker.vue
+++ b/docs/src/components/VersionPicker.vue
@@ -34,6 +34,8 @@ function absoluteUrl(relative) {
const props = defineProps<{ screenMenu?: boolean }>();
const versions = ref
holds a heading: keep them on one line */
+.vp-doc summary:has(> h3) {
+ cursor: pointer;
+ margin-top: 32px;
+}
+
+.vp-doc summary > h3 {
+ display: inline;
+ margin: 0;
+}
diff --git a/docs/src/_api_template.md b/docs/src/_api_template.md
index 28cfa3155..2e25d7e86 100644
--- a/docs/src/_api_template.md
+++ b/docs/src/_api_template.md
@@ -75,13 +75,15 @@ PARAMSKEY
heading_level: 3
show_root_full_path: false
-::: pysr.ParametricExpressionSpec
+::: pysr.AbstractExpressionSpec
options:
show_root_heading: true
heading_level: 3
show_root_full_path: false
-::: pysr.AbstractExpressionSpec
+## Type Specifications
+
+::: pysr.TypeSpec
options:
show_root_heading: true
heading_level: 3
@@ -100,3 +102,121 @@ PARAMSKEY
show_root_heading: true
heading_level: 3
show_root_full_path: false
+
+## Mutations
+
+::: pysr.ConstantMutation
+ options:
+ show_root_heading: true
+ heading_level: 3
+ show_root_full_path: false
+
+::: pysr.OperatorMutation
+ options:
+ show_root_heading: true
+ heading_level: 3
+ show_root_full_path: false
+
+::: pysr.FeatureMutation
+ options:
+ show_root_heading: true
+ heading_level: 3
+ show_root_full_path: false
+
+::: pysr.SwapOperandsMutation
+ options:
+ show_root_heading: true
+ heading_level: 3
+ show_root_full_path: false
+
+::: pysr.AddNodeMutation
+ options:
+ show_root_heading: true
+ heading_level: 3
+ show_root_full_path: false
+
+::: pysr.InsertNodeMutation
+ options:
+ show_root_heading: true
+ heading_level: 3
+ show_root_full_path: false
+
+::: pysr.DeleteNodeMutation
+ options:
+ show_root_heading: true
+ heading_level: 3
+ show_root_full_path: false
+
+::: pysr.RotateTreeMutation
+ options:
+ show_root_heading: true
+ heading_level: 3
+ show_root_full_path: false
+
+::: pysr.BacksolveMutation
+ options:
+ show_root_heading: true
+ heading_level: 3
+ show_root_full_path: false
+
+::: pysr.SimplifyMutation
+ options:
+ show_root_heading: true
+ heading_level: 3
+ show_root_full_path: false
+
+::: pysr.RandomizeMutation
+ options:
+ show_root_heading: true
+ heading_level: 3
+ show_root_full_path: false
+
+::: pysr.OptimizeMutation
+ options:
+ show_root_heading: true
+ heading_level: 3
+ show_root_full_path: false
+
+::: pysr.DoNothingMutation
+ options:
+ show_root_heading: true
+ heading_level: 3
+ show_root_full_path: false
+
+::: pysr.AbstractMutation
+ options:
+ show_root_heading: true
+ heading_level: 3
+ show_root_full_path: false
+
+## Plugins
+
+::: pysr.AdaptiveMutationWeightsPlugin
+ options:
+ show_root_heading: true
+ heading_level: 3
+ show_root_full_path: false
+
+::: pysr.AdaptiveParsimonyPlugin
+ options:
+ show_root_heading: true
+ heading_level: 3
+ show_root_full_path: false
+
+::: pysr.MutationBurstPlugin
+ options:
+ show_root_heading: true
+ heading_level: 3
+ show_root_full_path: false
+
+::: pysr.SimulatedAnnealingPlugin
+ options:
+ show_root_heading: true
+ heading_level: 3
+ show_root_full_path: false
+
+::: pysr.AbstractPlugin
+ options:
+ show_root_heading: true
+ heading_level: 3
+ show_root_full_path: false
diff --git a/docs/src/_index.md b/docs/src/_index.md
index 52587ba3b..df6abef05 100644
--- a/docs/src/_index.md
+++ b/docs/src/_index.md
@@ -16,7 +16,7 @@ hero:
text: View on GitHub
link: https://github.com/astroautomata/PySR
image:
- src: https://ai.damtp.cam.ac.uk/symbolicregression/dev/logo.png
+ src: https://raw.githubusercontent.com/MilesCranmer/PySR_Docs/master/images/pysr-logo.png
alt: PySR
features:
@@ -33,8 +33,8 @@ features:
details: "Customize everything: operators, loss functions, complexity, input types, optimizer, and more."
- icon: 🐍
- title: Python Native
- details: Built for Python with Julia backend for maximum performance and ease of use.
+ title: Familiar Python API
+ details: A scikit-learn interface, with results exported to SymPy, PyTorch, or JAX.
---
If you find PySR useful, please cite the paper [arXiv:2305.01582](https://arxiv.org/abs/2305.01582).
diff --git a/docs/src/backend.md b/docs/src/backend.md
index 801f353cb..3dc4b68cb 100644
--- a/docs/src/backend.md
+++ b/docs/src/backend.md
@@ -33,9 +33,9 @@ The main search code can be found in `src/SymbolicRegression.jl`.
Here are some tips:
-- The documentation for the backend is given [here](https://ai.damtp.cam.ac.uk/symbolicregression/dev/).
+- The documentation for the backend is given [here](https://julia.pysr.ai/dev/).
- Throughout the package, you will often see template functions which typically use a symbol `T` (such as in the string `where {T<:Real}`). Here, `T` is simply the datatype of the input data and stored constants, such as `Float32` or `Float64`. Writing functions in this way lets us write functions generic to types, while still having access to the specific type specified at compilation time.
-- Expressions are stored as binary trees, using the `Node{T}` type, described [here](https://ai.damtp.cam.ac.uk/symbolicregression/dev/types/#SymbolicRegression.CoreModule.EquationModule.Node).
+- Expressions are stored as binary trees, using the `Node{T}` type, described [here](https://julia.pysr.ai/dev/types#SymbolicRegression.CoreModule.EquationModule.Node).
- For reference, the main loop itself is found in the `equation_search` function inside [`src/SymbolicRegression.jl`](https://github.com/astroautomata/SymbolicRegression.jl/blob/master/src/SymbolicRegression.jl).
- Parts of the code which are typically edited by users include:
- [`src/CheckConstraints.jl`](https://github.com/astroautomata/SymbolicRegression.jl/blob/master/src/CheckConstraints.jl), particularly the function `check_constraints`. This function checks whether a given expression satisfies constraints, such as having a complexity lower than `maxsize`, and whether it contains any forbidden nestings of functions.
@@ -70,6 +70,6 @@ For more information on `juliapkg.json`, see [`pyjuliapkg`](https://github.com/J
## Additional notes
-If you get comfortable enough with the backend, you might consider using the Julia package directly: the API is given on the [SymbolicRegression.jl documentation](https://ai.damtp.cam.ac.uk/symbolicregression/dev/).
+If you get comfortable enough with the backend, you might consider using the Julia package directly: the API is given on the [SymbolicRegression.jl documentation](https://julia.pysr.ai/dev/).
If you make a change that you think could be useful to other users, don't hesitate to open a pull request on either the PySR or SymbolicRegression.jl repositories! Contributions are very appreciated.
diff --git a/docs/src/components/VersionPicker.vue b/docs/src/components/VersionPicker.vue
index 17514e206..2dc8a6b49 100644
--- a/docs/src/components/VersionPicker.vue
+++ b/docs/src/components/VersionPicker.vue
@@ -34,6 +34,8 @@ function absoluteUrl(relative) {
const props = defineProps<{ screenMenu?: boolean }>();
const versions = ref>([]);
const currentVersion = ref('Versions');
+// Releases are shown as vX.Y; links still point at the full vX.Y.Z folder.
+const minorLabel = (v: string) => /^(v\d+\.\d+)\.\d+$/.exec(v)?.[1] ?? v;
const isClient = ref(false);
const { site } = useData();
@@ -71,10 +73,21 @@ const loadVersions = async () => {
const scriptsLoaded = await waitForScriptsToLoad();
if (scriptsLoaded && window.DOC_VERSIONS && window.DOCUMENTER_CURRENT_VERSION) {
- versions.value = window.DOC_VERSIONS.map(v => ({
- text: v,
- link: absoluteUrl(`/${v}/`),
- }));
+ // DOC_VERSIONS lists releases newest first, so the first vX.Y.Z of each minor is its latest patch.
+ // Aliases (v1, v1.5) and prereleases stay reachable by URL, out of the picker.
+ const seenMinors = new Set();
+ versions.value = window.DOC_VERSIONS
+ .filter(v => {
+ if (v === 'dev' || v === 'stable') return true;
+ const minor = /^(v\d+\.\d+)\.\d+$/.exec(v)?.[1];
+ if (!minor || seenMinors.has(minor)) return false;
+ seenMinors.add(minor);
+ return true;
+ })
+ .map(v => ({
+ text: minorLabel(v),
+ link: absoluteUrl(`/${v}/`),
+ }));
currentVersion.value = window.DOCUMENTER_CURRENT_VERSION;
} else {
versions.value = [{ text: 'dev', link: absoluteUrl('/dev/') }];
@@ -89,12 +102,12 @@ const loadVersions = async () => {
isClient.value = true;
};
-const versionItems = computed(() => {
- return versions.value.map((v) => ({
- text: v.text,
- link: v.link
- }));
-});
+const versionItems = computed(() => versions.value.map((v) => ({
+ text: v.text,
+ link: v.link,
+ target: '_self',
+ noIcon: true
+})));
onMounted(() => {
if (typeof window !== 'undefined') {
@@ -108,12 +121,12 @@ onMounted(() => {
Data generation code
+
+```python
+import numpy as np
+
+ALPHABET = "ABCDEFGHIJKLMNOPQRSTUVWXYZ"
+PLAIN = "THELAWSOFPHYSICSAREWRITTENINTHELANGUAGEOFMATHEMATICS"
+CIPHER = "".join(ALPHABET[(5 * ALPHABET.index(c) + 8) % 26] for c in PLAIN)
+
+X = np.array([[ALPHABET.index(c)] for c in CIPHER], dtype=object)
+y = np.array([ALPHABET.index(c) for c in PLAIN], dtype=object)
+```
+
+
+
+`Letter` stores a code in `0:25`. Its `string` hook displays constants as letters, such as `V` for code `21`.
+
+```python
+from pysr import PySRRegressor, TypeSpec
+
+spec = TypeSpec(
+ "Letter",
+ fields={"code": "Int"},
+ sample="rng -> Letter(rand(rng, 0:25))",
+ mutate="(rng, value, temperature) -> Letter(mod(value.code + rand(rng, (-1, 1)), 26))",
+ string="value -> string(Char('A' + value.code))",
+)
+```
+
+
+The three operators implement addition, subtraction, and multiplication modulo $26$. Every result remains a `Letter`, and the modulus appears directly in each operator. That built-in alphabet size makes this the easier cipher search.
+
+Use absolute code distance as `elementwise_loss` to reward partial improvements. Zero loss means every ciphertext/plaintext pair matches:
+
+```python
+from pysr import PySRRegressor
+
+model = PySRRegressor(
+ type_spec=spec,
+ operators={
+ 2: [
+ "shift(a::Letter, b::Letter) = Letter(mod(a.code + b.code, 26))",
+ "unshift(a::Letter, b::Letter) = Letter(mod(a.code - b.code, 26))",
+ "mix(a::Letter, b::Letter) = Letter(mod(a.code * b.code, 26))",
+ ]
+ },
+ elementwise_loss="letter_loss(prediction::Letter, target::Letter)::Float64 = abs(prediction.code - target.code)",
+ niterations=20,
+ random_state=0,
+)
+model.fit(X, y)
+print(model.equations_[["complexity", "loss", "equation"]].to_string(index=False))
+```
+
+
+A `loss_function` is useful when scoring needs the dataset as a whole, as in the rollout in [Swinging up a cart-pole](/examples/objectives#swinging-up-a-cart-pole-with-a-rollout-objective) or the period certificate in [Inventing a pseudorandom generator](/examples/objectives#inventing-a-pseudorandom-generator-with-no-target).
+
+One recovered expression is `mix(V, shift(S, x0))`: `V` has code $21$ and `S` has code $18$, giving $21(x_0 + 18) \bmod 26$. Since $21$ is the multiplicative inverse of $5$ modulo $26$, this expression decodes the ciphertext.
+
+### Learning the modulus too
+
+The harder variant removes the alphabet size from the operators. It uses values in `0:51` with integer addition, multiplication, and modulo, so the search must discover the constant $26$ along with the decryption law. This removes a strong structural hint and enlarges the constant search.
+
+One recovered decryption law is:
+
+```
+modulo(mul(add(18, x0), 47), 26)
+```
+
+Run `examples/affine_cipher.py` for the first search or `examples/affine_freemod.py` for the full type, operator, and loss setup of this variant. Allow several minutes for the latter.
+
+## Rediscovering Conway's Game of Life
+
+
+
+Conway's Game of Life can be written as a rule about a cell and its live-neighbour count. This example withholds that rule and supplies only its transition table, asking PySR to reconstruct the logical expression.
+
+The table has $18$ distinct states: two current cell states, each paired with every neighbour count from $0$ through $8$. Repeating the table $8$ times produces $144$ rows for scoring without introducing new cases. A zero-one loss therefore charges $1/18$ for a rule that misses one distinct state. The arrays use `dtype=object` so each value can be wrapped in the custom Julia type.
+
+
+Data generation code
+
+```python
+import numpy as np
+
+TABLE = [(a, n, int(n == 3 or (a == 1 and n == 2))) for a in (0, 1) for n in range(9)]
+REPLICAS = 8
+X = np.array([[a, n] for a, n, _ in TABLE] * REPLICAS, dtype=object)
+y = np.array([t for _, _, t in TABLE] * REPLICAS, dtype=object)
+```
+
+
+
+The `Cell` type wraps a single `Int`, which is enough for the current state, the neighbour count, and the output code. Its hooks propose small counts, move existing constants locally modulo $9$, and print the integer payload.
+
+```python
+from pysr import PySRRegressor, TypeSpec
+
+spec = TypeSpec(
+ "Cell",
+ fields={"v": "Int"},
+ sample="rng -> Cell(rand(rng, 0:8))",
+ mutate="(rng, value, temperature) -> Cell(mod(value.v + rand(rng, (-1, 1)), 9))",
+ string="value -> string(value.v)",
+)
+```
+
+
+The operator vocabulary contains `not`, `and`, `or`, and the equality test `eq`. Arithmetic and threshold comparisons are absent, so a neighbour threshold can enter an expression only as a constant in a call such as `eq(n, 3)`. The loss is zero-one equality on each table row. It records exact agreement, which suits a discrete truth table where the numerical distance between cell codes carries no meaning.
+
+```python
+from pysr import PySRRegressor
+
+model = PySRRegressor(
+ type_spec=spec,
+ operators={
+ 1: ["not(a::Cell) = Cell(a.v == 0 ? 1 : 0)"],
+ 2: [
+ "and(a::Cell, b::Cell) = Cell((a.v != 0 && b.v != 0) ? 1 : 0)",
+ "or(a::Cell, b::Cell) = Cell((a.v != 0 || b.v != 0) ? 1 : 0)",
+ "eq(a::Cell, b::Cell) = Cell(a.v == b.v ? 1 : 0)",
+ ],
+ },
+ elementwise_loss="cell_loss(prediction::Cell, target::Cell)::Float64 = prediction.v == target.v ? 0.0 : 1.0",
+ niterations=160,
+ random_state=0,
+)
+model.fit(X, y, variable_names=["alive", "n"])
+print(model.equations_[["complexity", "loss", "equation"]].to_string(index=False))
+```
+
+
+A recovered rule is:
+
+```
+or(eq(n, 3), and(eq(n, 2), alive))
+```
+
+The first branch allows birth or survival at three neighbours. The second adds survival of a live cell at two neighbours. It agrees with all 18 possible input states, so the transition table fully checks this local rule.
+
+Run `examples/game_of_life.py` to reproduce the Game of Life search.
+
+## Searching over machine words
+
+
+
+This search targets the bit interleave used by Morton, or Z-order, indexing. It gives every expression a `Word` value and restricts the vocabulary to machine instructions: xor, and, or, not, left shift, and right shift. Shift distances are ordinary `Word` constants, so the search must find them as part of the expression.
+
+The reference interleave is evaluated on all $256$ pairs of $4$-bit inputs.
+
+Data generation code
+
+```python
+import numpy as np
+
+from pysr import PySRRegressor, TypeSpec
+
+WIDTH = 4
+WIDER = 5
+
+
+def morton(x, y, width=WIDTH):
+ z = 0
+ for i in range(width):
+ z |= ((x >> i) & 1) << (2 * i)
+ z |= ((y >> i) & 1) << (2 * i + 1)
+ return z
+
+
+PAIRS = [(a, b) for b in range(1 << WIDTH) for a in range(1 << WIDTH)]
+X = np.array(PAIRS, dtype=object)
+y = np.array([morton(a, b) for a, b in PAIRS], dtype=object)
+```
+
+
+
+The `Word` type wraps a `UInt32`. Its `sample` hook favors small values because useful shift distances are small while mask constants still need to reach across the full word. `mutate` either flips a single bit or takes a random unsigned step. This lets a mask change one bit at a time. The `string` hook renders small values as decimals and larger values as zero-padded hexadecimal, keeping shift constants and masks easy to distinguish. The `shl` and `shr` definitions below pass `b.bits` directly; this example does not mask the shift count into `0:31`, so arbitrary sampled `UInt32` counts follow the underlying Julia shift semantics.
+
+```python
+from pysr import TypeSpec
+
+spec = TypeSpec(
+ "Word",
+ fields={"bits": "UInt32"},
+ sample="rng -> Word(rand(rng, Bool) ? rand(rng, UInt32(0):UInt32(31)) : rand(rng, UInt32))",
+ mutate="(rng, value, temperature) -> Word(rand(rng, Bool) ? xor(value.bits, UInt32(1) << rand(rng, 0:31)) : value.bits + rand(rng, (UInt32(1), typemax(UInt32))))",
+ string='value -> value.bits < UInt32(32) ? string(value.bits) : "0x" * string(value.bits, base = 16, pad = 8)',
+)
+```
+
+
+Each operator lifts one machine instruction to `Word`. The loss counts the set bits in the xor of prediction and target and divides by $32$. A zero loss therefore means that every output bit agrees for every scored pair.
+
+```python
+from pysr import PySRRegressor
+
+model = PySRRegressor(
+ type_spec=spec,
+ operators={
+ 1: ["bnot(a::Word) = Word(~a.bits)"],
+ 2: [
+ "bxor(a::Word, b::Word) = Word(xor(a.bits, b.bits))",
+ "band(a::Word, b::Word) = Word(a.bits & b.bits)",
+ "bor(a::Word, b::Word) = Word(a.bits | b.bits)",
+ "shl(a::Word, b::Word) = Word(a.bits << b.bits)",
+ "shr(a::Word, b::Word) = Word(a.bits >> b.bits)",
+ ],
+ },
+ elementwise_loss="bit_loss(prediction::Word, target::Word)::Float64 = count_ones(xor(prediction.bits, target.bits)) / 32",
+ maxsize=45,
+ niterations=800,
+ random_state=0,
+ verbosity=0,
+)
+
+model.fit(X, y, variable_names=["x", "y"])
+```
+
+
+Set `maxsize` large enough for the intended computation. A four-bit reference interleaver uses 41 tree nodes because repeated subexpressions each count toward its size; `maxsize=45` leaves room for that construction.
+
+An expression that matches all four-bit pairs need not work at larger widths: the training domain leaves parts of the masks unconstrained. Use the script's `exact_front_rows` helper to check four-bit agreement and `wider_agreement(model, index)` to test a selected expression on five-bit inputs.
+
+Run `examples/morton_interleave.py` to reproduce this machine-word search.
+
+## Turtle graphics: searching over drawings
+
+
+
+The excerpts below come from `examples/turtle_graphics.py`. Run that script for the complete example; it supplies the Python `path` converter and the Julia `PREAMBLE` containing the underscore-prefixed drawing, mutation, and scoring helpers and constants.
+
+Here the value is a drawing program. A `Path` is a list of pen commands, `forward` and `turn`; operators repeat a path, reflect it, and concatenate command lists. Each expression therefore evaluates to a program, and the value of that program is the picture it draws.
+
+The drawing uses two parallel vectors. `kind=0` marks forward motion. `kind=1` marks a turn. `val` stores the forward length or the turn angle in radians. The `ARM` list below writes turn angles in degrees, and the Python `path` helper converts those entries to radians before storing `y`; the custom `string` hook prints stored turns back in degrees. The target is a plus-sign outline made from four copies of an arm with three edges. The dataset has one row and one feature, `f`, containing a unit forward step. Turns and the path operators must supply the rest of the program.
+
+```python
+import numpy as np
+from pysr import PySRRegressor, TypeSpec
+
+ARM = [("F", 1.0), ("T", 90.0), ("F", 1.0), ("T", -90.0), ("F", 1.0), ("T", 90.0)]
+TARGET = ARM * 4
+
+X = np.empty((1, 1), dtype=object)
+X[0, 0] = ([0], [1.0])
+y = np.empty(1, dtype=object)
+y[0] = path(TARGET)
+VARIABLE_NAMES = ["f"]
+```
+
+
+### Scoring a picture
+
+Two programs can trace the same outline in different orders, so the loss compares rasterized pictures instead of command lists. Each drawing is placed on a $96\times96$ grid after its bounding box is centered and its larger extent is scaled to the grid. The comparison therefore responds to shape while ignoring translation and overall size.
+
+The score is a symmetric chamfer distance between the two sets of lit pixels, measured in pixels and divided by the grid width. The score averages the distance from candidate pixels to the target and from target pixels to the candidate, so zero means the raster masks coincide.
+
+The chamfer loss and the raster
+
+```julia
+# Symmetric chamfer distance in pixels, divided by the grid width. Zero means
+# the two drawings lit exactly the same pixels.
+function _chamfer(pk::Vector{Int8}, pv::Vector{Float64},
+ tk::Vector{Int8}, tv::Vector{Float64})
+ a = _raster(pk, pv)
+ a === nothing && return _FAIL
+ b, db = _target_grids(tk, tv)
+ any(b) || return _FAIL
+ da = _edt(a)
+ ab = _mean_to(a, db); ba = _mean_to(b, da)
+ (ab === nothing || ba === nothing) && return _FAIL
+ v = 0.5 * (ab + ba) / _GRID
+ return isfinite(v) ? v : _FAIL
+end
+```
+
+
+
+### The pen program type
+
+The `TypeSpec` for `Path` stores the two vectors and supplies hooks for continuous turns. `sample` creates one turn at an angle in the interval from $-180$ to $180$ degrees. `mutate` selects a turn and either perturbs it by a Gaussian step with scale $30$ degrees times the temperature or resamples it. The scalar-constant hooks expose every turn to BFGS, which can polish an angle after evolution has placed it near a good value.
+
+Constant folding can collapse an all-constant subtree into a multi-command path. `mutate` and `string` therefore handle values command by command, and the printer represents a folded multi-command constant as `p[...]`.
+
+```python
+from pysr import TypeSpec
+
+PATH = TypeSpec(
+ "Path",
+ fields={"kind": "Vector{Int8}", "val": "Vector{Float64}"},
+ preamble=PREAMBLE,
+ sample="rng -> Path(Int8[1], [_real(rng)])",
+ mutate="""(rng, value, temperature) -> begin
+ turns = findall(isone, value.kind)
+ isempty(turns) && return Path(Int8[1], [_real(rng)])
+ i = rand(rng, turns)
+ val = copy(value.val)
+ val[i] = rand(rng, Bool) ? _perturb(rng, val[i], temperature) : _real(rng)
+ Path(copy(value.kind), val)
+ end""",
+ scalar_constants="value -> _turns(value.kind, value.val)",
+ with_scalar_constants="(value, c) -> Path(copy(value.kind), _put_turns(value.kind, value.val, c))",
+ is_valid="value -> !isempty(value.kind) && all(isfinite, value.val)",
+ string="value -> _literal(value.kind, value.val)",
+)
+```
+
+
+The five operators form a small program algebra. `seq` runs one path after another. `dup` and `tri` repeat a path twice and three times. `mir` negates every turn, reflecting the path. `nest` performs an L-system substitution: it replaces each forward command with a copy of the full path. Each copy is scaled so its net displacement matches the forward command it replaces. Repetition and reflection let a short expression draw many strokes.
+
+```python
+OPERATORS = {
+ 1: [
+ "dup(a::Path) = (t = _seq(a.kind, a.val, a.kind, a.val); Path(t[1], t[2]))",
+ "tri(a::Path) = (t = _seq(a.kind, a.val, a.kind, a.val); u = _seq(t[1], t[2], a.kind, a.val); Path(u[1], u[2]))",
+ "mir(a::Path) = (t = _mir(a.kind, a.val); Path(t[1], t[2]))",
+ "nest(a::Path) = (t = _nest(a.kind, a.val); Path(t[1], t[2]))",
+ ],
+ 2: [
+ "seq(a::Path, b::Path) = (t = _seq(a.kind, a.val, b.kind, b.val); Path(t[1], t[2]))"
+ ],
+}
+
+CHAMFER_LOSS = "chamfer(prediction::Path, target::Path)::Float64 = _chamfer(prediction.kind, prediction.val, target.kind, target.val)"
+```
+
+
+### The search
+
+Pass the path type, operators, and picture loss to the estimator:
+
+```python
+from pysr import PySRRegressor
+
+model = PySRRegressor(
+ type_spec=PATH,
+ operators=OPERATORS,
+ elementwise_loss=CHAMFER_LOSS,
+ niterations=20,
+ verbosity=0,
+ random_state=0,
+)
+
+model.fit(X, y, variable_names=["f"])
+print(model.equations_[["complexity", "loss", "equation"]].to_string(index=False))
+```
+
+
+Scoring each candidate requires rasterization and a distance transform, so even this 20-iteration search can take tens of minutes to hours. Zero loss means the raster masks match; it does not require identical command order or exact turn angles.
+
+One recovered drawing program is:
+
+```
+dup(dup(mir(seq(t-179.6493032061754, tri(seq(p[t-270.04570985370526,t-270.04570985370526,t-270.04570985370526], f))))))
+```
+
+The complete drawing search is in `examples/turtle_graphics.py`.
+
+## Evolving a hopping controller
+
+
+
+Each value in this search is a motor command with three numbers, one for each joint of MuJoCo's hopper. The objective measures performance in simulation. Each candidate expression controls the robot for a fixed-horizon rollout, and its loss is the negated reward from that rollout. The search uses no expert policy, demonstration data, or distilled network; it optimises locomotion directly.
+
+The snippets below describe the recorded search. The runnable script, `examples/hopper_controller.py`, replays its saved controller.
+
+The controller is a `Vec3`, three `Float64` components carried through the tree together:
+
+```python
+from pysr import TypeSpec
+
+TypeSpec(
+ "Vec3",
+ fields={"data": "Vector{Float64}"},
+ sample="rng -> Vec3(randn(rng, 3))",
+ scalar_constants="value -> Float64[value.data[1], value.data[2], value.data[3]]",
+ with_scalar_constants="(value, constants) -> Vec3(collect(constants))",
+ is_valid="value -> length(value.data) == 3 && all(isfinite, value.data)",
+ string='value -> "Vec3([" * join(repr.(Float64.(value.data)), ", ") * "])"',
+ loss_type="Float64",
+)
+```
+
+Each sampled constant contains three independent Gaussian draws, giving each
+motor its own bias. The scalar-constant hooks expose all three components to the
+optimiser so it can tune each motor separately.
+
+The operators are componentwise addition, subtraction and multiplication, and nothing else. No control-specific primitive was supplied:
+
+```python
+{
+ 2: [
+ 'vadd(x, y) = Vec3(x.data .+ y.data)',
+ 'vsub(x, y) = Vec3(x.data .- y.data)',
+ 'vmul(x, y) = Vec3(x.data .* y.data)',
+ ]
+}
+```
+
+### Scoring a controller by simulating it
+
+The eleven features encode the hopper state as vectors. They include the three
+joint angles, angular velocities clipped to $[-10, 10]$ and divided by $10$, and two cyclic permutations
+of each. The torso height, pitch, forward velocity, vertical velocity and
+angular velocity are each broadcast across all three components. The three torso velocities are divided by $5$. A single tree
+computes all three motor commands, so the permuted copies allow each joint to use
+other joints' angles. Before reaching the motors, the final value passes through
+a componentwise `tanh` that limits each command to the actuator range.
+
+Each rollout accumulates the standard Hopper-v5 reward, a survival term plus forward velocity minus a small control cost, and divides by the full horizon, so a fall contributes zero for every step it did not survive. That single number, averaged over eight frozen training starts and negated, is the loss. The Julia loss function is a thin call into the rollout evaluator:
+
+```python
+LOSS_SOURCE = "hopper_rollout_loss(tree, dataset, options) = Main.HopperObjective.loss(tree, dataset, options)"
+```
+
+Because the objective never reads the dataset, `X` and `y` are one row of zero vectors and batching is off. All the information comes from the simulator.
+
+### The recovered controller
+
+Here `A` denotes joint angles, `V` scaled joint angular velocities, and `P` the cyclic permutation `(x0, x1, x2) -> (x1, x2, x0)`. The scalars `p`, `f`, `u`, and `w` denote torso pitch and scaled forward, vertical, and angular velocities:
+
+```
+D = A - 2 P(A) + P(V) - 2 (f + w) + c0
+L = -A - (p + w (P(V) - w)) D - c1
+R = -P(A) - u + w + c2
+action = tanh(P(A) - w + L R)
+```
+
+All products are componentwise, with scalars broadcast. The constant vectors are `c0 = (-0.2477, -4.2268, 8.8397)`, `c1 = (0.0645, 0.6181, 0.3969)`, and `c2 = (0.4974, -1.1921, 1.1808)`, rounded here; the replay script retains their full precision.
+
+### Reproducing the result
+
+Replay the controller on twenty held-out starts with twice the training reset noise.
+
+The replay checks whether the hopper stays upright for eight seconds, travels five metres, and completes at least four contact/liftoff cycles. These criteria check sustained hopping beyond the reward used during training.
+
+The result covers small perturbations of one starting pose in simulation. The evaluation does not test obstacles, hardware, or indefinite hopping.
+
+The saved controller was evaluated with Gymnasium 1.2.2 and MuJoCo 3.3.7. Install `gymnasium[mujoco]==1.2.2` and `mujoco==3.3.7`, then run `python examples/hopper_controller.py`.
diff --git a/docs/src/examples/expression-specifications.md b/docs/src/examples/expression-specifications.md
new file mode 100644
index 000000000..0f240d73a
--- /dev/null
+++ b/docs/src/examples/expression-specifications.md
@@ -0,0 +1,476 @@
+# Expression specifications
+
+## Expression specifications
+
+Use an expression specification when part of the equation is known in advance.
+It fixes the structure around learnable subexpressions while preserving PySR's ordinary
+fitting and prediction interfaces. Inspect template components through `model.equations_`;
+the arbitrary Julia `combine` string does not provide the usual SymPy or LaTeX exports.
+Choose `TemplateExpressionSpec` when the outer form is known and the inner functions remain unknown.
+
+### Template Expressions
+
+Write the outer form in `combine` and name each learnable hole in `expressions`.
+`variable_names` supplies the column names available to `combine`, and each
+hole receives the arguments written in its call. Here `f` receives `x_1` and `x_2`,
+while `g` receives `x_3`, so the search keeps those input roles separate. The target
+form is:
+
+$$ y = \sin(f(x_1, x_2)) + g(x_3) $$
+
+The code below creates the data, declares the template, and fits the model:
+
+```python
+import numpy as np
+
+from pysr import PySRRegressor, TemplateExpressionSpec
+
+# Create data
+X = np.random.randn(1000, 3)
+y = np.sin(X[:, 0] + X[:, 1]) + X[:, 2]**2
+
+# Define template: we want sin(f(x1, x2)) + g(x3)
+template = TemplateExpressionSpec(
+ expressions=["f", "g"],
+ variable_names=["x1", "x2", "x3"],
+ combine="sin(f(x1, x2)) + g(x3)",
+)
+
+model = PySRRegressor(
+ expression_spec=template,
+ binary_operators=["+", "*", "-", "/"],
+ unary_operators=["sin"],
+ maxsize=10,
+)
+model.fit(X, y)
+```
+
+### Parametric Expressions
+
+Add learnable parameter vectors to the template when categories share a formula
+but require different coefficients. This example uses category-specific scale and offset
+values for the target:
+
+$$ y = \alpha \sin(x_1) + \beta $$
+
+Here $\alpha$ is the category-specific scale and $\beta$ is the offset. The
+example has $3$ categories, with $\alpha \in \{1.0, 2.0, 0.5\}$ and
+$\beta \in \{0.1, 1.5, -0.5\}$.
+
+```python
+import numpy as np
+
+from pysr import PySRRegressor, TemplateExpressionSpec
+
+# Create data with 2 features and 3 categories
+X = np.random.uniform(-3, 3, (1000, 2))
+category = np.random.randint(0, 3, 1000)
+
+# Parameters for each category
+offsets = [0.1, 1.5, -0.5]
+scales = [1.0, 2.0, 0.5]
+
+# y = scale[category] * sin(x1) + offset[category]
+y = np.array([
+ scales[c] * np.sin(x1) + offsets[c]
+ for x1, c in zip(X[:, 0], category)
+])
+```
+
+Declare the parameter-vector lengths, optimize their entries during the search,
+and use the selected values in the template. The following block defines that
+parametric expression:
+
+```python
+from pysr import TemplateExpressionSpec
+
+template = TemplateExpressionSpec(
+ expressions=["f"],
+ variable_names=["x1", "x2", "category"],
+ parameters={"p1": 3, "p2": 3}, # One parameter per category
+ combine="f(x1, x2, p1[category], p2[category])"
+)
+```
+
+Append the category as the column listed in `variable_names`; the template indexes
+that column when it selects parameter entries. Python labels start at zero, whereas
+Julia arrays start at one, so shift the labels before stacking them.
+
+```python
+import numpy as np
+
+category_p_one = category + 1
+X_with_category = np.column_stack([X, category_p_one])
+```
+
+Fit the model with the augmented feature matrix, and retain this category column
+for any later `predict` call. The template uses it to choose the corresponding
+parameter entry:
+
+```python
+from pysr import PySRRegressor
+
+model = PySRRegressor(
+ expression_spec=template,
+ binary_operators=["+", "*", "-", "/"],
+ unary_operators=["sin"],
+ maxsize=10,
+)
+model.fit(X_with_category, y)
+
+# Evaluate the fitted model on the same feature/category rows:
+model.predict(X_with_category)
+```
+
+See [Expression Specifications](/api#expression-specifications) for the complete
+reference to this API. The same pattern extends to several template holes and to parameter
+vectors whose entries vary by category.
+
+### Learning multiple outputs jointly
+
+`TemplateExpressionSpec` can couple several scalar holes through one scalar
+residual. Put the observed components in `X`, let the template construct each
+prediction, and return the combined error. This is useful when the components share
+a known outer term.
+
+The target consists of $3$ components. Each contains the shared term $\exp(x_1)$
+and its own second term:
+
+$$\begin{align*}
+y_1 &= \exp(x_1) + x_2^2 \\
+y_2 &= \exp(x_1) + \sin(x_3) \\
+y_3 &= \exp(x_1) + x_1 \cdot x_2
+\end{align*}$$
+
+The data and their noise are generated in the following block:
+
+```python
+import numpy as np
+
+from pysr import PySRRegressor, TemplateExpressionSpec
+
+n = 200
+rstate = np.random.RandomState(0)
+x1 = rstate.uniform(-2, 2, n)
+x2 = rstate.uniform(-2, 2, n)
+x3 = rstate.uniform(-2, 2, n)
+
+# True model with shared component exp(x1):
+y1 = np.exp(x1) + x2**2
+y2 = np.exp(x1) + np.sin(x3)
+y3 = np.exp(x1) + x1 * x2
+
+# Add some noise
+y1 += 0.05 * rstate.randn(n)
+y2 += 0.05 * rstate.randn(n)
+y3 += 0.05 * rstate.randn(n)
+```
+
+`X` contains the three inputs followed by the three observed target components,
+in the same order listed in `variable_names`. The template can therefore use the
+target columns while building its residual:
+
+```python
+import numpy as np
+
+X = np.column_stack([x1, x2, x3, y1, y2, y3])
+```
+
+The template returns the sum of squared errors across the three components for each row.
+
+These sums constrain the predicted components, while the individual `shared` and `f*` terms are not uniquely identifiable. A common function can be added to `shared` and subtracted from every `f*` without changing the residual, so the values in the evaluation block are conditional on recovering the intended decomposition.
+
+```python
+from pysr import TemplateExpressionSpec
+
+spec = TemplateExpressionSpec(
+ expressions=["f1", "f2", "f3", "shared"],
+ variable_names=["x1", "x2", "x3", "y1", "y2", "y3"],
+ combine="""
+ v = shared(x1, x2, x3)
+ y1_predicted = v + f1(x1, x2, x3)
+ y2_predicted = v + f2(x1, x2, x3)
+ y3_predicted = v + f3(x1, x2, x3)
+
+ residuals = (
+ abs2(y1 - y1_predicted) +
+ abs2(y2 - y2_predicted) +
+ abs2(y3 - y3_predicted)
+ )
+
+ residuals
+ """
+)
+```
+
+The fit uses a dummy target and an `elementwise_loss` that returns the template
+value, because `combine` already returns a row-wise squared residual. The estimator
+still receives `dummy_y` to satisfy the fit interface; those zeros do not enter the
+residual:
+
+```python
+import numpy as np
+
+from pysr import PySRRegressor
+
+model = PySRRegressor(
+ expression_spec=spec,
+ binary_operators=["+", "-", "*", "/"],
+ unary_operators=["exp", "sin"],
+ maxsize=20,
+ niterations=50,
+ elementwise_loss="(pred, target) -> pred",
+)
+
+dummy_y = np.zeros(n)
+model.fit(X, dummy_y)
+```
+
+Inspect the fitted named expressions through their Julia objects:
+
+```python
+# Simply get the expression with the highest score:
+idx = model.equations_.score.idxmax()
+
+# Extract the Julia object:
+julia_expr = model.equations_.loc[idx, 'julia_expression']
+
+# Access individual subexpressions:
+for name in ['f1', 'f2', 'f3', 'shared']:
+ tree = getattr(julia_expr.trees, name)
+ print(f"{name}: {tree}")
+```
+
+The named trees can also be evaluated directly through the Julia bridge:
+
+```python
+import numpy as np
+
+from pysr import jl
+from pysr.julia_helpers import jl_array
+
+SR = jl.SymbolicRegression
+
+# Get individual trees
+f1_tree = julia_expr.trees.f1
+shared_tree = julia_expr.trees.shared
+
+# Evaluate one point (x1=1, x2=2, x3=3). `eval_tree_array` expects
+# features by rows, so this is a 3-by-1 array.
+test_inputs = jl_array(np.array([[1.0], [2.0], [3.0]]))
+f1_result, _ = SR.eval_tree_array(f1_tree, test_inputs, model.julia_options_)
+shared_result, _ = SR.eval_tree_array(shared_tree, test_inputs, model.julia_options_)
+
+print(f"f1 at (1,2,3): {f1_result[0]}") # Intended target term: x2^2 = 4.0
+print(f"shared at (1,2,3): {shared_result[0]}") # Intended shared term: exp(1) ≈ 2.718
+```
+
+## Recovering a magnetic field from force measurements
+
+
+
+The force law is fixed in this example, while its coefficient functions are unknown.
+A charged particle in a viscous medium follows
+
+$$ F = -\eta(T)\,v + v \times B(t), $$
+
+The data provide $(t, v, T)$ and the resulting force. The template fixes the drag and
+cross-product assembly and leaves, as learnable holes, the three components of $B(t)$
+and the temperature-dependent drag scale $\eta(T)$.
+
+Each target is a three-component value, so the example defines a `Force` type with three
+fields. Pairing `TypeSpec` with `TemplateExpressionSpec` lets `PySRRegressor` pass these
+custom values through the fit while the template constructs each force from the known law.
+Here the template reads the `.x` field from each returned hole; the other fields remain
+part of the `Force` representation.
+
+The complete runnable script is `examples/magnetic_field.py`; the blocks below expose its
+data and model definitions in the same order.
+
+
+Data generation code
+
+Each scalar input is embedded as a diagonal `Force`, with the same value in each of the
+three fields. The output remains a full three-component force. This lets the template recover
+scalar inputs from `.x` while keeping one custom value type for both `X` and `y`:
+
+`TypeSpec` converts each tuple in these object arrays into the generated `Force` using the
+declared field order (`x`, `y`, `z`), so Python does not need to call `Force(...)` directly.
+
+```python
+import numpy as np
+
+from pysr import PySRRegressor, TemplateExpressionSpec, TypeSpec
+
+OMEGA = 2 * np.pi
+
+
+def experiments(n, seed):
+ rng = np.random.default_rng(seed)
+ T = 298.15 + 0.5 * rng.random(n)
+ t = 10.0 * rng.random(n)
+ v = 2.0 * rng.random((n, 3)) - 1.0
+ B = np.stack([np.sin(OMEGA * t), np.cos(OMEGA * t), np.exp(-t / 10.0)], axis=1)
+ F = -1e-5 * np.sqrt(T)[:, None] * v + np.cross(v, B)
+ return np.column_stack([t, v[:, 0], v[:, 1], v[:, 2], T]), F
+
+
+def as_forces(rows):
+ """An (n, 3) float array as a column of n Force values."""
+ out = np.empty(len(rows), dtype=object)
+ out[:] = [tuple(map(float, row)) for row in rows]
+ return out
+
+
+def on_diagonal(columns):
+ out = np.empty(columns.shape, dtype=object)
+ for j in range(columns.shape[1]):
+ out[:, j] = as_forces(np.repeat(columns[:, j, None], 3, axis=1))
+ return out
+
+
+INPUTS, FORCES = experiments(1000, seed=0)
+X = on_diagonal(INPUTS)
+y = as_forces(FORCES)
+variable_names = ["t", "v_x", "v_y", "v_z", "T"]
+
+HELD_OUT_INPUTS, HELD_OUT_FORCES = experiments(500, seed=12345)
+HELD_OUT_X = on_diagonal(HELD_OUT_INPUTS)
+```
+
+
+
+`FORCE` declares the three fields and the hooks that control sampling, constant
+optimization, string output, and failed evaluation. `sample` draws a value for each
+field. The paired `scalar_constants` and `with_scalar_constants` hooks expose and restore
+all three components, so every `Force` constant is a vector. `init_invalid` supplies a
+`Force` with `NaN` in every field. No explicit `is_valid` is given, so the
+default finite-constant check rejects that value:
+
+```python
+from pysr import TypeSpec
+
+FORCE = TypeSpec(
+ "Force",
+ fields={"x": "Float64", "y": "Float64", "z": "Float64"},
+ sample="rng -> Force(randn(rng, 3)...)",
+ definitions="""
+ Base.sin(a::Force)::Force = Force(sin(a.x), sin(a.y), sin(a.z))
+ Base.cos(a::Force)::Force = Force(cos(a.x), cos(a.y), cos(a.z))
+ function Base.sqrt(a::Force)::Force
+ f(v) = v < 0 ? NaN : sqrt(v)
+ return Force(f(a.x), f(a.y), f(a.z))
+ end
+ Base.exp(a::Force)::Force = Force(exp(a.x), exp(a.y), exp(a.z))
+ Base.:+(a::Force, b::Force)::Force = Force(a.x + b.x, a.y + b.y, a.z + b.z)
+ Base.:-(a::Force, b::Force)::Force = Force(a.x - b.x, a.y - b.y, a.z - b.z)
+ Base.:*(a::Force, b::Force)::Force = Force(a.x * b.x, a.y * b.y, a.z * b.z)
+ Base.:/(a::Force, b::Force)::Force = Force(a.x / b.x, a.y / b.y, a.z / b.z)
+ """,
+ scalar_constants="value -> [value.x, value.y, value.z]",
+ with_scalar_constants="(value, c) -> Force(c[1], c[2], c[3])",
+ string="value -> sprint(show, (value.x, value.y, value.z); context = :compact => true)",
+ init_invalid="() -> Force(NaN, NaN, NaN)",
+)
+```
+
+`definitions` runs after the `Force` type is generated, so its Julia methods can return
+`Force` values. The guarded `sqrt` returns `NaN` for negative fields instead of throwing.
+Register the methods by arity and define a loss over all three force components:
+
+```python
+OPERATORS = {
+ 1: [
+ "Base.sin",
+ "Base.cos",
+ "Base.sqrt",
+ "Base.exp",
+ ],
+ 2: [
+ "Base.:+",
+ "Base.:-",
+ "Base.:*",
+ "Base.:/",
+ ],
+}
+
+FORCE_LOSS = "force_loss(a::Force, b::Force)::Float64 = (a.x - b.x)^2 + (a.y - b.y)^2 + (a.z - b.z)^2"
+```
+
+Inside a template `combine`, PySR's backend wraps input columns and named-hole results in
+`ValidVector` values. The wrapper exposes `.x` for raw values and `.valid` for evaluation
+status; `ValidVector(raw, valid)` constructs a wrapper for a new result. This type is supplied
+by the template runtime rather than defined in this example.
+
+`PHYSICS` checks the hole results for validity, then assembles one `Force` per row
+using the fixed cross product and drag law:
+
+```python
+from pysr import TemplateExpressionSpec
+
+PHYSICS = r"""
+begin
+ _B_x = B_x(t)
+ _B_y = B_y(t)
+ _B_z = B_z(t)
+ _F_d_scale = F_d_scale(T)
+ if !(_B_x.valid && _B_y.valid && _B_z.valid && _F_d_scale.valid)
+ return ValidVector(_B_x.x, false)
+ end
+ F = map(_B_x.x, _B_y.x, _B_z.x, _F_d_scale.x,
+ v_x.x, v_y.x, v_z.x) do bx, by, bz, fd, ux, uy, uz
+ b1, b2, b3 = bx.x, by.x, bz.x
+ u1, u2, u3 = ux.x, uy.x, uz.x
+ s = fd.x
+ Force(u2 * b3 - u3 * b2 + s * u1,
+ u3 * b1 - u1 * b3 + s * u2,
+ u1 * b2 - u2 * b1 + s * u3)
+ end
+ ValidVector(F, true)
+end
+"""
+
+STRUCTURE = TemplateExpressionSpec(
+ combine=PHYSICS,
+ expressions=["B_x", "B_y", "B_z", "F_d_scale"],
+ variable_names=variable_names,
+)
+```
+
+The field holes receive only `t`; `F_d_scale` receives only `T`. Only the `.x` fields
+of their `Force` values enter the numerical residual. The `.y` and `.z` fields can still
+affect whether intermediate values remain valid.
+
+```python
+from pysr import PySRRegressor
+
+model = PySRRegressor(
+ type_spec=FORCE,
+ expression_spec=STRUCTURE,
+ operators=OPERATORS,
+ elementwise_loss=FORCE_LOSS,
+ niterations=100,
+ verbosity=0,
+ random_state=0,
+)
+
+model.fit(X, y, variable_names=variable_names)
+print(model.equations_[["complexity", "loss", "equation"]].to_string(index=False))
+```
+
+An example fitted expression is:
+
+ B_x = sin(#1 / (0.159155, 6.43746, -0.55788))
+ B_y = cos(#1 / (0.159155, 6.43746, -0.55788))
+ B_z = exp(#1 * (-0.1, -0.334652, -1.00446))
+ F_d_scale = -0.000172745, 0.262294, 0.471951
+
+`#1` is the first argument passed to each hole: `t` for the field components and
+`T` for the drag scale. Constants print three entries because they are full `Force`
+values. The field expressions match the sinusoidal and exponential forms used to
+generate the data; the constant drag scale only approximates the temperature-dependent law.
+
+The script validates force predictions on independently generated experiments with
+`held_out_error(model)`. Accurate force predictions alone do not establish exact
+recovery of every coefficient function. Expect an expensive search: a run took about
+1.5 hours in the recorded full-vector configuration.
diff --git a/docs/src/examples/getting-started.md b/docs/src/examples/getting-started.md
new file mode 100644
index 000000000..5b380d4cd
--- /dev/null
+++ b/docs/src/examples/getting-started.md
@@ -0,0 +1,153 @@
+# Getting started
+
+## Simple search
+
+This target combines a cosine term, a square term, and a constant.
+
+The model searches with `+`, `-`, `*`, and `/`. To include the cosine used in the target, add `unary_operators=["cos"]`.
+
+```python
+import numpy as np
+
+from pysr import PySRRegressor
+
+X = 2 * np.random.randn(100, 5)
+y = 2 * np.cos(X[:, 3]) + X[:, 0] ** 2 - 2
+model = PySRRegressor(binary_operators=["+", "-", "*", "/"])
+model.fit(X, y)
+print(model)
+```
+
+## Custom operator
+
+A custom operator has a definition for the symbolic search and a callable for SymPy export and for evaluating equations. Here `inv(x) = 1/x` defines the Julia-side reciprocal operator, while `extra_sympy_mappings` gives that name its Python implementation.
+
+```python
+import numpy as np
+
+from pysr import PySRRegressor
+
+X = 2 * np.random.randn(100, 5)
+y = 1 / X[:, 0]
+model = PySRRegressor(
+ binary_operators=["+", "*"],
+ unary_operators=["inv(x) = 1/x"],
+ extra_sympy_mappings={"inv": lambda x: 1/x},
+)
+model.fit(X, y)
+print(model)
+```
+
+The generic `operators` parameter replaces the separate `unary_operators` and `binary_operators`. Its integer keys specify arity, so a dictionary can describe unary, binary, and higher-arity operators. See the [Operators of any arity](/examples/search-behaviour#operators-of-any-arity) section for this form.
+
+## Multiple outputs
+
+A two-dimensional target makes the model multi-output. Here each column of `y` is the reciprocal of a different input feature, so the fitted model supplies one selected equation for each output.
+
+```python
+import numpy as np
+
+from pysr import PySRRegressor
+
+X = 2 * np.random.randn(100, 5)
+y = 1 / X[:, [0, 1, 2]]
+model = PySRRegressor(
+ binary_operators=["+", "*"],
+ unary_operators=["inv(x) = 1/x"],
+ extra_sympy_mappings={"inv": lambda x: 1/x},
+)
+model.fit(X, y)
+```
+
+## Plotting an expression
+
+After fitting the multi-output model, `model.latex()[0]` returns the LaTeX form of the selected equation for output 0.
+
+```python
+model.latex()[0]
+```
+
+Use `model.latex()[1]` for output 1.
+
+The plot below compares the target and prediction for output 0.
+
+```python
+from matplotlib import pyplot as plt
+plt.scatter(y[:, 0], model.predict(X)[:, 0])
+plt.xlabel('Truth')
+plt.ylabel('Prediction')
+plt.show()
+```
+
+
+
+The `index` argument chooses a candidate equation from `model.equations_`. Pass it to `predict`, `sympy`, or `latex` to inspect or evaluate another candidate. For multiple outputs, provide one equation index per output in output order, then select the desired output column from the predictions.
+
+## Feature selection
+
+Feature selection reduces the candidate input columns before symbolic regression. Set `select_k_features=5` to use a random-forest preprocessor that retains up to five candidates. For structured high-dimensional data, the paper [2006.11287](https://arxiv.org/abs/2006.11287) discusses breaking the problem into smaller pieces before applying PySR.
+
+The target in this example uses features 3 and 19.
+
+```python
+import numpy as np
+
+X = np.random.randn(300, 30)
+y = X[:, 3]**2 - X[:, 19]**2 + 1.5
+```
+
+Configure feature selection in the model:
+
+```python
+from pysr import PySRRegressor
+
+model = PySRRegressor(
+ binary_operators=["+", "-", "*", "/"],
+ unary_operators=["exp"],
+ select_k_features=5,
+)
+```
+
+Fit the model to select the input columns before the search:
+
+```python
+model.fit(X, y)
+```
+
+PySR prints the selected feature names before starting the search, for example:
+
+```text
+Using features ['x3', 'x5', 'x7', 'x19', 'x21']
+```
+
+The names use zero-based column indices. This selection includes both target features, `x3` and `x19`, plus three other candidates. The symbolic search uses only these retained columns.
+
+## Denoising
+
+When per-observation weights are available, pass `weights` to `fit` and define `elementwise_loss` with the matching signature. Without weights, the function accepts `(prediction, target)`. With weights, it accepts `(prediction, target, weight)`. The weighted example uses `elementwise_loss="myloss(x, y, w) = w * (x - y)^2"`. Pass a full custom objective as `loss_function`.
+
+The `denoise=True` option preprocesses the targets with a Gaussian process before symbolic regression. Its predictions become the targets searched by PySR. Because this example passes no `Xresampled`, the predictions are made at the original rows in `X`.
+
+The target combines an exponential term, two linear terms, and random noise:
+
+```python
+import numpy as np
+
+X = np.random.randn(100, 5)
+noise = np.random.randn(100) * 0.1
+y = np.exp(X[:, 0]) + X[:, 1] + X[:, 2] + noise
+```
+
+Create and fit the denoising model:
+
+```python
+from pysr import PySRRegressor
+
+model = PySRRegressor(
+ binary_operators=["+", "-", "*", "/"],
+ unary_operators=["exp"],
+ denoise=True,
+)
+model.fit(X, y)
+print(model)
+```
diff --git a/docs/src/examples/instrumentation.md b/docs/src/examples/instrumentation.md
new file mode 100644
index 000000000..035440646
--- /dev/null
+++ b/docs/src/examples/instrumentation.md
@@ -0,0 +1,227 @@
+# Instrumentation and workflow
+
+## Using TensorBoard for logging
+
+Pass `TensorBoardLoggerSpec` through `logger_spec` to write search progress to TensorBoard. `TensorBoardLoggerSpec` records the model's search hyperparameters and writes scalar summaries for the run. The example uses `min_loss`, which tracks the current best loss, and `pareto_volume`, which summarizes the Pareto front.
+
+```python
+import numpy as np
+
+from pysr import PySRRegressor, TensorBoardLoggerSpec
+
+rstate = np.random.RandomState(42)
+
+# Uniform dist between -3 and 3:
+X = rstate.uniform(-3, 3, (1000, 2))
+y = np.exp(X[:, 0]) + X[:, 1]
+
+# Create a logger that writes to "logs/run*":
+logger_spec = TensorBoardLoggerSpec(
+ log_dir="logs/run",
+ log_interval=10, # Log every 10 iterations
+)
+
+model = PySRRegressor(
+ binary_operators=["+", "*", "-", "/"],
+ logger_spec=logger_spec,
+)
+model.fit(X, y)
+```
+
+Start TensorBoard with the directory that contains the run logs:
+
+```bash
+tensorboard --logdir logs/
+```
+
+## Recording the genealogy of a search
+
+
+
+The complete recorder and post-processing workflow is in `examples/search_trace.py`. `load_trace` and `ancestry` are example-local helpers, not PySR API methods.
+
+Tracing records how a symbolic-regression population changes. Setting `use_tracing=True` activates the optional recorder. PySR loads the required JSON dependency lazily on first use. It then writes JSONL records to `tracing_file`, one for each population at each iteration, with the live members and their mutation, crossover, tuning, and death events. These recorded fields supply the members and parent relationships that the helpers need to reconstruct a graph and one selected path. The helpers work from recorded events and do not expose every historical branch.
+
+With tracing disabled, PySR does not request any trace output or load the tracing dependency.
+
+The target combines one quadratic term, one cosine term, and an offset. The added Gaussian noise has standard deviation equal to $5\%$ of the target's spread:
+
+```python
+import numpy as np
+
+rng = np.random.default_rng(0)
+X = rng.uniform(-3, 3, size=(200, 3))
+y = 2.5382 * np.cos(X[:, 2]) + X[:, 0] ** 2 - 1.5
+y = y + rng.normal(0, 0.05 * y.std(), size=200)
+```
+
+```python
+from pysr import PySRRegressor
+
+model = PySRRegressor(
+ binary_operators=["+", "-", "*", "/"],
+ unary_operators=["cos", "exp"],
+ population_size=27,
+ populations=8,
+ ncycles_per_iteration=90,
+ maxsize=20,
+ niterations=25,
+ random_state=0,
+ use_tracing=True,
+ tracing_file="search_trace.jsonl",
+)
+model.fit(X, y)
+```
+
+Keep tracing runs short: the files grow with the search, and this `load_trace` helper reads the entire file into memory.
+
+
+Trace-reading code
+
+```python
+import json
+import pathlib
+
+def load_trace(trace_path):
+ nodes, edges = {}, []
+ for line in pathlib.Path(trace_path).read_text().splitlines():
+ record = json.loads(line)
+ if record["record_type"] != "iteration":
+ continue
+ for member in record["members"]:
+ nodes.setdefault(member["ref"], {}).update(member)
+ for ref, entry in record["mutations"].items():
+ parent = int(ref)
+ nodes.setdefault(parent, {}).update(
+ {k: entry[k] for k in ("tree", "loss", "cost", "parent") if k in entry}
+ )
+ for event in entry["events"]:
+ kind = event["type"]
+ detail = event.get("mutation") or event.get("details") or {}
+ if kind == "crossover" and event["parent1"] != parent:
+ continue
+ children = (
+ (event["child1"], event["child2"])
+ if kind == "crossover"
+ else (event["child"],) if "child" in event else ()
+ )
+ for child in children:
+ edges.append(
+ dict(
+ parent=parent,
+ child=child,
+ kind=kind,
+ mutation=detail.get("type"),
+ result=detail.get("result", "accept"),
+ mate=event.get("parent2"),
+ )
+ )
+ for edge in edges:
+ if (
+ edge["mate"] is not None
+ and nodes.get(edge["child"], {}).get("parent") == edge["mate"]
+ ):
+ edge["parent"], edge["mate"] = edge["mate"], edge["parent"]
+ return nodes, edges
+```
+
+
+
+`ancestry` selects the lowest-loss recorded member and follows the first edge for each child. It returns one path, which ends when no edge remains or a reference repeats.
+
+```python
+def ancestry(nodes, edges):
+ first_edge = {}
+ for edge in edges:
+ first_edge.setdefault(edge["child"], edge)
+ winner = min(
+ (ref for ref, n in nodes.items() if n.get("loss") is not None),
+ key=lambda ref: nodes[ref]["loss"],
+ )
+ chain, current, seen = [], winner, set()
+ while current in first_edge and current not in seen:
+ seen.add(current)
+ chain.append(first_edge[current])
+ current = first_edge[current]["parent"]
+ chain.reverse()
+ return winner, chain
+```
+
+Trace records label features `x1`, `x2`, and so on by column order, independently of any `variable_names` passed to the search. Thus `x1` refers to `X[:, 0]`, and `x3` refers to `X[:, 2]` in this example.
+
+Use `examples/search_trace.py` for the complete tracing workflow.
+
+## Closing an agent loop with `guesses=`
+
+The `guesses=` keyword accepts a list of expression strings to seed a fit. An external workflow can inspect `model.equations_`, propose an expression, and pass it into a later search. PySR can then optimize its numeric constants.
+
+This example replays a recorded coding-agent proposal from `PROPOSALS`. The custom `agent(front, tried)` helper returns stored expressions; it does not call a live agent or model service. The complete implementation is in `examples/agent_loop_guesses.py`.
+
+The target is relativistic kinetic energy in natural units, with $c=1$ and speeds up to $0.88c$. Gaussian noise has standard deviation $0.15\%$ of the median exact energy. `NOISE_FLOOR` stores its variance for the loop's training-loss stopping rule.
+
+
+Data generation code
+
+```python
+import numpy as np
+
+rng = np.random.default_rng(42)
+mass = rng.uniform(0.5, 3.0, 160)
+velocity = rng.uniform(0.08, 0.88, 160)
+kinetic_energy_exact = mass * (1 / np.sqrt(1 - velocity**2) - 1)
+NOISE_SIGMA = 0.0015 * float(np.median(kinetic_energy_exact))
+kinetic_energy = kinetic_energy_exact + rng.normal(0.0, NOISE_SIGMA, mass.size)
+NOISE_FLOOR = NOISE_SIGMA**2
+
+X = np.stack([mass, velocity], axis=1)
+y = kinetic_energy
+VARIABLE_NAMES = ["mass", "velocity"]
+```
+
+
+
+```python
+from pysr import PySRRegressor
+
+model_kwargs = dict(
+ operators={1: ["sqrt"], 2: ["+", "-", "*", "/"]},
+ niterations=10,
+ verbosity=0,
+)
+```
+
+The proposal uses the known relativistic structure:
+
+```python
+PROPOSALS = ["(mass / sqrt(1 - (velocity * velocity))) - mass"]
+
+def agent(front, tried):
+ print("front handed to the agent:")
+ print(
+ front[["complexity", "loss", "equation"]].to_string(index=False),
+ "\n",
+ flush=True,
+ )
+ return PROPOSALS[len(tried)] if len(tried) < len(PROPOSALS) else None
+```
+
+The loop begins with `guesses=None` and passes each proposal as a one-element list to the next fit. It stops when the minimum training loss reaches `2 * NOISE_FLOOR` or the helper has no proposal left. This threshold is an example-specific stopping rule, not a lower bound on achievable loss.
+
+```python
+from pysr import PySRRegressor
+
+guesses, tried = None, []
+while True:
+ model = PySRRegressor(**model_kwargs, guesses=guesses, random_state=0)
+ model.fit(X, y, variable_names=VARIABLE_NAMES)
+ front = model.equations_
+ if front["loss"].min() <= 2 * NOISE_FLOOR:
+ break
+ proposal = agent(front, tried)
+ if proposal is None:
+ break
+ tried.append(proposal)
+ guesses = [proposal]
+```
+
+Use `examples/agent_loop_guesses.py` for the complete handoff workflow.
diff --git a/docs/src/examples/objectives.md b/docs/src/examples/objectives.md
new file mode 100644
index 000000000..1766b33c8
--- /dev/null
+++ b/docs/src/examples/objectives.md
@@ -0,0 +1,332 @@
+# Objectives and losses
+
+## Custom objectives
+
+`loss_function` gives the objective the candidate's expression tree. Use it when the score needs to manipulate the tree or read dataset-level information. Use `loss_function_expression` when the score needs the complete expression object. Both modes accept Julia source with one of the signatures below.
+
+```julia
+objective(tree_or_expression, dataset, options)
+objective(tree_or_expression, dataset, options, idx=nothing)
+```
+
+With three arguments, the objective sees the dataset given to the current evaluation. Automatic batching makes that dataset the active batch. A method with `idx=nothing` can receive the full dataset and the selected row indices, allowing the objective to slice the rows itself.
+
+The following objective reproduces ordinary unweighted mean squared error, while retaining access to the optional batch indices:
+
+```python
+from pysr import PySRRegressor
+
+objective = """
+function mse_objective(tree, dataset::Dataset{T,L}, options, idx=nothing)::L where {T,L}
+ X = idx === nothing ? dataset.X : dataset.X[:, idx]
+ y = idx === nothing ? dataset.y : dataset.y[idx]
+ prediction, complete = eval_tree_array(tree, X, options)
+ complete || return L(Inf)
+ return sum(abs2, prediction .- y) / length(y)
+end
+"""
+
+model = PySRRegressor(
+ loss_function=objective,
+ binary_operators=["+", "-", "*", "/"],
+)
+```
+
+`eval_tree_array` returns a completion flag with the predictions. Check that flag before using the array and assign an infinite, or suitably large, loss to an incomplete evaluation. Return the dataset's loss type `L`; it must remain real even when the data type `T` is complex.
+
+PySR imports `SymbolicRegression` into Julia before evaluating these source strings, making `Dataset` and `eval_tree_array` available without another import.
+
+A full objective can also attach a different meaning to the stored tree. Here the two children of a binary root supply the numerator and denominator, forming the rational function $P(X)/Q(X)$. The fully qualified `DE.get_child` calls below access the backend's tree interface.
+
+```python
+from pysr import PySRRegressor
+
+objective = """
+function rational_objective(tree, dataset::Dataset{T,L}, options, idx=nothing)::L where {T,L}
+ tree.degree == 2 || return L(Inf)
+ X = idx === nothing ? dataset.X : dataset.X[:, idx]
+ y = idx === nothing ? dataset.y : dataset.y[idx]
+
+ numerator = SymbolicRegression.InterfaceDynamicExpressionsModule.DE.get_child(tree, 1)
+ denominator = SymbolicRegression.InterfaceDynamicExpressionsModule.DE.get_child(tree, 2)
+ p, p_complete = eval_tree_array(numerator, X, options)
+ q, q_complete = eval_tree_array(denominator, X, options)
+ p_complete && q_complete || return L(Inf)
+
+ prediction = p ./ q
+ return sum(abs2, prediction .- y) / length(y)
+end
+"""
+
+model = PySRRegressor(
+ loss_function=objective,
+ binary_operators=["+", "-"],
+)
+```
+
+This objective ignores the root operator and interprets only its two children. The equation table displays the stored tree. `predict` and symbolic export likewise retain the raw tree and cannot automatically reconstruct the quotient. Keep the objective source and apply the child transformation yourself whenever you evaluate the selected equation.
+
+## Writing the objective in Python
+
+The objective itself can live in Python while Julia supplies a thin shim. The search releases Python's global interpreter lock (GIL) while it runs, so the shim reacquires the lock, calls the Python function, and converts its scalar result to the dataset's loss type. The Python function receives the tree, dataset, and options as Julia objects, and it can call back into Julia, including `eval_tree_array`:
+
+```python
+import numpy as np
+
+from pysr import PySRRegressor, jl
+
+
+def python_objective(tree, dataset, options):
+ prediction, completed = jl.SymbolicRegression.eval_tree_array(
+ tree, dataset.X, options
+ )
+ if not completed:
+ return float("inf")
+
+ prediction = np.asarray(prediction)
+ target = np.asarray(dataset.y)
+ return float(np.mean((prediction - target) ** 2))
+
+
+jl.python_objective = python_objective
+
+jl.seval("""
+using PythonCall
+
+function python_objective_shim(
+ tree, dataset::Dataset{T,L}, options
+)::L where {T,L}
+ PythonCall.GIL.@lock begin
+ return pyconvert(L, python_objective(tree, dataset, options))
+ end
+end
+""")
+
+model = PySRRegressor(
+ loss_function="python_objective_shim",
+ precision=64,
+)
+```
+
+`precision=64` makes the data and loss use `Float64`, matching the Python float returned by this objective.
+
+Use this bridge when the objective needs Python-specific logic. Every evaluation crosses the Julia/Python boundary, and concurrent callbacks serialize on the GIL. Keep throughput-sensitive objectives in Julia.
+
+## Swinging up a cart-pole with a rollout objective
+
+
+
+A supervised example gives a target for each row. A control policy has no label of that kind for its force: a force is good only if the subsequent trajectory is good. `loss_function_expression` provides the candidate expression to a rollout objective that can evaluate it throughout the closed loop. This example runs 500 steps from each of 16 training starts and returns the negative mean reward per step. Most of the starts leave the pole hanging down, so the expression must swing it up, catch it, and balance it with one closed-form policy.
+
+### The plant and what the policy sees
+
+The implemented plant uses a $1.0\,\mathrm{kg}$ cart and a $0.1\,\mathrm{kg}$ pole with half-length $0.5\,\mathrm{m}$, gravity $9.8$, and semi-implicit Euler integration with $dt=0.02\,\mathrm{s}$. Its 500-step horizon is therefore $10\,\mathrm{s}$. The policy output is clipped to `[-1, 1]` and multiplied by $10\,\mathrm{N}$. The five inputs are scaled cart position, scaled cart speed, $\sin\theta$, $\cos\theta$, and scaled pole rate. The sine/cosine pair avoids the wrap discontinuity at $\pm\pi$, while the cosine input `c` supplies a direct uprightness signal.
+
+
+Constants and the plant step
+
+```python
+import numpy as np
+
+M_C, M_P, LENGTH, GRAVITY = 1.0, 0.1, 0.5, 9.8
+FORCE_CAP, DT = 10.0, 0.02
+RAIL, V_SCALE, OMEGA_SCALE = 2.4, 4.0, 8.0
+HORIZON = 500 # 10 s of control at dt = 0.02
+
+
+def step(state, force):
+ x, v, theta, omega = state.T
+ total = M_C + M_P
+ sin, cos = np.sin(theta), np.cos(theta)
+ q = (force + M_P * LENGTH * omega**2 * sin) / total
+ theta_dd = (GRAVITY * sin - cos * q) / (LENGTH * (4 / 3 - M_P * cos**2 / total))
+ x_dd = q - M_P * LENGTH * theta_dd * cos / total
+ v, omega = v + DT * x_dd, omega + DT * theta_dd
+ return np.stack([x + DT * v, v, theta + DT * omega, omega], axis=-1)
+```
+
+
+
+Use the initial-state arrays and `observe` helper in `examples/cartpole_objective.py` to construct `X`. Its columns contain the five policy observations, and `y` is a zero placeholder that the objective ignores. The 16 training starts cover hanging, horizontal, and near-upright states; the 64 held-out starts use wider ranges of position, speed, angle, and pole rate.
+
+### The reward
+
+At each step, the objective accumulates the reward
+
+$$ r = 2\cos\theta - 0.05\,x^2 - 0.01\left(\frac{F}{10}\right)^2 - 0.005\left(\frac{F - F_{\text{prev}}}{10}\right)^2 $$
+
+The upright term contributes up to 2 per step. The remaining terms penalize, in order, cart displacement, normalized actuator force, and normalized force changes. Here $F$ is the clipped force and $F_{\mathrm{prev}}$ is the preceding force. The objective returns the negative mean of $r$ over 500 steps and 16 starts, so an upright stationary policy has loss near $-2$.
+
+The complete objective source is defined in `examples/cartpole_objective.py` as the Python string `CARTPOLE_REWARD`. The block below shows its Julia function body and abbreviates the state unpacking, observation-buffer setup, integration step, and divergence check supplied by that script. Treat it as a behavior excerpt, not a standalone definition.
+
+```julia
+function cartpole_reward(ex, dataset::Dataset{T,L}, options)::L where {T,L}
+ n = size(dataset.X, 2)
+ # ... unpack x, v, theta, omega from the observation columns ...
+ for _ in 1:500
+ # ... refill obs from the current state ...
+ raw, ok = eval_tree_array(ex, obs, options)
+ (ok && length(raw) == n && all(isfinite, raw)) || return L(Inf)
+
+ force = 10.0 .* clamp.(raw, -1.0, 1.0)
+ earned += sum(
+ 2.0 .* cos_t .- 0.05 .* x .^ 2 .-
+ 0.01 .* (force ./ 10.0) .^ 2 .-
+ 0.005 .* ((force .- previous) ./ 10.0) .^ 2
+ )
+ previous = force
+ # ... one semi-implicit Euler step, then the divergence check ...
+ end
+ return L(-earned / (500 * n))
+end
+```
+
+Invalid evaluations and divergent trajectories receive a penalty.
+
+### The search
+
+```python
+import sympy
+
+from pysr import PySRRegressor
+
+model = PySRRegressor(
+ operators={
+ 1: ["square", "abs", "tanh"],
+ 2: ["+", "-", "*", "/", "max", "min"],
+ 3: ["ifelse(t, a, b) = t > 0 ? a : b"],
+ },
+ extra_sympy_mappings={
+ "ifelse": lambda t, a, b: sympy.Piecewise((a, t > 0), (b, True))
+ },
+ loss_function_expression=CARTPOLE_REWARD,
+ loss_scale="linear", # rewards make the loss negative, which log scaling forbids
+ maxsize=28,
+ maxdepth=10,
+ parsimony=0.001,
+ niterations=120,
+ populations=32,
+ population_size=64,
+ random_state=0,
+)
+model.fit(X, y, variable_names=["x_n", "v_n", "s", "c", "omega_n"])
+```
+
+`loss_scale="linear"` is required here because logarithmic scaling accepts nonnegative losses, including zero, while this reward can make the loss negative. The operator set supplies arithmetic, `max`, and `min`; the user-defined ternary `ifelse` can switch between swing-up and balancing behavior. `extra_sympy_mappings` translates that custom operator to a SymPy `Piecewise` expression for export. The search uses multithreading, so repeated fits can differ even with the same nominal seed.
+
+### Evaluating a policy
+
+One controller returned by the search is:
+
+```
+(tanh(((0.46341985 * v_n) + (omega_n + (s + omega_n))) / 0.00043365502) - v_n) / 0.97654253
+```
+
+A low training loss alone does not establish successful swing-up from new initial states. In `examples/cartpole_objective.py`, `mean_rewards(model, HELD_OUT_STARTS, index=...)` evaluates a selected equation on the wider held-out set, and `check(model)` requires positive mean reward on every start. Use these rollout checks when choosing a policy.
+
+Allow roughly two hours for the search configuration above. The complete runnable example is `examples/cartpole_objective.py`.
+
+## Inventing a pseudorandom generator with no target
+
+
+
+Ordinary regression supplies `y` and compares every candidate output with a target. This example instead searches for a 32-bit state update $x \mapsto f(x)$ whose iterates have specified properties. Its score comes from how the map behaves, so the objective defines what counts as a good generator without consulting a target array.
+
+The data rows provide 32 random nonzero seeds. `y` is a zero column retained for the estimator interface; the objective never reads it.
+
+```python
+import numpy as np
+
+from pysr import PySRRegressor, TypeSpec
+
+WIDTH = 32
+SEEDS = [int(v) for v in np.random.default_rng(0).integers(1, 1 << WIDTH, size=32)]
+X = np.empty((len(SEEDS), 1), dtype=object)
+for i, seed in enumerate(SEEDS):
+ X[i, 0] = seed
+y = np.zeros(len(SEEDS), dtype=object)
+```
+
+Each state is an unsigned 32-bit word, wrapped in a Julia type. The `shl` and `shr` operators below pass their `UInt32` shift count directly; this example does not mask counts into `0:31`, so arbitrary sampled counts follow Julia's shift semantics and bits shifted out of the 32-bit word are discarded. The operations do not rotate the state. The `sample` hook favors values from 0 through 31 because those are common shift distances, while still sampling arbitrary `UInt32` values. Half of mutations flip a single bit, allowing a local change to a tap constant; the other half adds either 1 or `typemax(UInt32)`. The `string` hook displays values below 32 in decimal and larger words as eight hexadecimal digits.
+
+```python
+from pysr import TypeSpec
+
+SPEC = TypeSpec(
+ "Word",
+ fields={"bits": "UInt32"},
+ sample="rng -> Word(rand(rng, Bool) ? rand(rng, UInt32(0):UInt32(31)) : rand(rng, UInt32))",
+ mutate="(rng, value, temperature) -> Word(rand(rng, Bool) ? xor(value.bits, UInt32(1) << rand(rng, 0:31)) : value.bits + rand(rng, (UInt32(1), typemax(UInt32))))",
+ string='value -> value.bits < UInt32(32) ? string(value.bits) : "0x" * string(value.bits, base = 16, pad = 8)',
+ loss_type="Float64",
+)
+```
+
+The operators form a restricted bitwise instruction set: complement, XOR, AND, OR, left shift, and right shift. The search excludes arithmetic and floating-point operations.
+
+### What the objective scores
+
+`examples/prng_period.py` defines the objective as a Julia source string named `PRNG_LOSS`. It includes the period-certificate implementation used below.
+
+For every candidate, the objective walks 256 steps from each seed and adds six terms that penalize departures from their desired properties:
+
+- `period` is $1 - \log_2(\mathrm{proven})/32$, with `proven` equal to the certified orbit length when certification succeeds.
+- `balance` is the mean squared bias of each bit column away from one half.
+- `autocorr` is the mean squared normalized self-correlation of each bit column at lags 1, 2, 3, and 5.
+- `diffusion` penalizes an average of fewer than four changed output bits when each input bit is flipped on a fixed probe set.
+- `edge` penalizes output bits that respond to fewer than two distinct input-bit flips across that probe set.
+- `shear` measures correlations between bit `b` and shifted bit `b + d` at one- and two-step lags, for nonzero shifts `d` from -3 to 3. It keeps a plain shift from receiving a good score.
+
+### Certifying a period of four billion without walking it
+
+The requested nonzero orbit contains $2^{32}-1 = 4{,}294{,}967{,}295$ states, so a direct walk is impractical. The certificate evaluates the candidate at three groups of inputs: zero, all 32 basis states, and 32 fixed probes. It requires zero to map to zero and the image of each probe to agree with the matrix assembled from the basis images. These finite checks only screen for GF(2) linearity on the supplied evidence; they do not prove global linearity for an arbitrary expression.
+
+A successful certificate gives the candidate's period provided the expression is globally GF(2)-linear. If certification fails, the objective uses the orbit length observed during the finite walk. The period-checking implementation is in `examples/prng_period.py`.
+
+### Settings
+
+The settings below pass the source-defined `PRNG_LOSS` objective to PySR:
+
+```python
+from pysr import PySRRegressor
+
+model = PySRRegressor(
+ type_spec=SPEC,
+ operators={
+ 1: ["bnot(a::Word) = Word(~a.bits)"],
+ 2: [
+ "bxor(a::Word, b::Word) = Word(xor(a.bits, b.bits))",
+ "band(a::Word, b::Word) = Word(a.bits & b.bits)",
+ "bor(a::Word, b::Word) = Word(a.bits | b.bits)",
+ "shl(a::Word, b::Word) = Word(a.bits << b.bits)",
+ "shr(a::Word, b::Word) = Word(a.bits >> b.bits)",
+ ],
+ },
+ loss_function=PRNG_LOSS,
+ batching=False,
+ maxsize=20,
+ populations=8,
+ population_size=30,
+ ncycles_per_iteration=30,
+ niterations=800,
+ random_state=0,
+ verbosity=0,
+)
+model.fit(X, y, variable_names=["x"])
+```
+
+Keep `batching=False` because the objective walks every seed; a minibatch would score a different problem. Each candidate requires a 256-step walk over all 32 seeds. Allow about an hour for this search.
+
+### Results
+
+The script's `check(model)` looks for a front member that passes a separate finite linearity screen and whose inferred GF(2) matrix has order $2^{32}-1$. As with the objective's certificate, this establishes the candidate's period only if the expression is globally GF(2)-linear.
+
+One recovered expression is:
+
+```
+bxor(shl(x, 1), bxor(shr(bxor(shr(x, 1), x), 1), band(x, 0xffffffe3)))
+```
+
+Fixed shifts, XOR, and AND with a constant mask make this expression globally linear by construction. For such a map, matrix order $2^{32}-1$ implies that every nonzero word lies on one orbit of that length. The remaining score terms assess the generator's finite-sample balance, correlations, and diffusion.
+
+The complete runnable example is `examples/prng_period.py`.
diff --git a/docs/src/examples/physics.md b/docs/src/examples/physics.md
new file mode 100644
index 000000000..156c66f27
--- /dev/null
+++ b/docs/src/examples/physics.md
@@ -0,0 +1,183 @@
+# Physics and units
+
+## Dimensional constraints
+
+Dimensional constraints let PySR use physical units during symbolic search. When the units of the features and target are known, pass a unit expression for each feature and one for the target so candidates with mismatched dimensions receive a penalty.
+
+The fit interface follows the unit syntax of `DynamicQuantities.jl`. The following example uses Astropy to sample a mass $M$, a test mass $m$, and a separation $r$, then evaluates Newton's inverse-square law $F = GMm/r^2$.
+
+```python
+import numpy as np
+from astropy import units as u, constants as const
+
+M = (np.random.rand(100) + 0.1) * const.M_sun
+m = 100 * (np.random.rand(100) + 0.1) * u.kg
+r = (np.random.rand(100) + 0.1) * const.R_earth
+G = const.G
+
+F = G * M * m / r**2
+```
+
+The force values span a wide range. This loss compares their magnitudes in logarithmic space and penalizes a wrong sign:
+
+```python
+elementwise_loss = """function loss_fnc(prediction, target)
+ scatter_loss = abs(log((abs(prediction)+1f-20) / (abs(target)+1f-20)))
+ sign_loss = 10 * (sign(prediction) - sign(target))^2
+ return scatter_loss + sign_loss
+end
+"""
+```
+
+`dimensional_constraint_penalty` adds a cost when a candidate fails dimensional analysis:
+
+```python
+from pysr import PySRRegressor
+
+model = PySRRegressor(
+ binary_operators=["+", "-", "*", "/"],
+ unary_operators=["square"],
+ elementwise_loss=elementwise_loss,
+ complexity_of_constants=2,
+ maxsize=25,
+ niterations=100,
+ populations=50,
+ # Amount to penalize dimensional violations:
+ dimensional_constraint_penalty=10**5,
+)
+```
+
+Pass numeric values to `fit`, with matching units in column order through `X_units` and the target unit through `y_units`. Unit strings follow the Julia syntax documented by [DynamicQuantities.jl](https://ai.damtp.cam.ac.uk/dynamicquantities/dev/#Usage).
+
+```python
+import pandas as pd
+
+# Get numerical arrays to fit:
+X = pd.DataFrame(dict(
+ M=M.to("M_sun").value,
+ m=m.to("kg").value,
+ r=r.to("R_earth").value,
+))
+y = F.value
+
+model.fit(
+ X,
+ y,
+ X_units=["Constants.M_sun", "kg", "Constants.R_earth"],
+ y_units="kg * m / s^2"
+)
+```
+
+The loss is nonnegative, and the dimensional penalty is additive. A candidate with total loss below that penalty therefore did not receive the dimensional-violation charge and is dimensionally consistent for the supplied units. The marker `"[⋅]"` denotes a constant whose units remain free. The unit-aware form below illustrates this notation: `M[kg]` supplies mass, so the constant supplies the acceleration units needed for the target:
+
+```julia
+"y[m s⁻² kg] = (M[kg] * 2.6353e-22[⋅])"
+```
+
+Set `dimensionless_constants_only=True` to restrict fitted constants to dimensionless quantities.
+
+## Using differential operators
+
+A [`TemplateExpressionSpec`](/examples/expression-specifications) can include differential operators in its `combine` string. `D` takes the expression to differentiate and an argument index. Here, `D(f, 1)` differentiates the learned function `f` with respect to its first argument, and `df(x)` evaluates that derivative. Matching those values to data turns antiderivative discovery into regression.
+
+The regression target is the integrand $1/(x^2\sqrt{x^2 - 1})$ on $x > 1$. The template searches for a function $f$ whose derivative matches samples of the integrand, leaving an additive constant undetermined.
+
+```python
+import numpy as np
+
+from pysr import PySRRegressor, TemplateExpressionSpec
+
+x = np.random.uniform(1, 10, (1000,)) # Integrand sampling points
+y = 1 / (x**2 * np.sqrt(x**2 - 1)) # Evaluation of the integrand
+
+expression_spec = TemplateExpressionSpec(
+ expressions=["f"],
+ variable_names=["x"],
+ combine="df = D(f, 1); df(x)",
+)
+
+model = PySRRegressor(
+ binary_operators=["+", "-", "*", "/"],
+ unary_operators=["sqrt"],
+ expression_spec=expression_spec,
+ maxsize=20,
+)
+model.fit(x[:, np.newaxis], y)
+```
+
+The target antiderivative is $f(x) = \frac{\sqrt{x^2 - 1}}{x}$, up to an additive constant.
+
+## Discovering a PDE
+
+For a field $u(x,t)$ sampled on a space-time grid, each point becomes one regression row: its field value and spatial derivatives are features, and its time derivative is the target. PySR then searches for the right-hand side in $u_t = f(u, u_x, u_{xx}, \ldots)$.
+
+The following example simulates the viscous Burgers equation $u_t = -u\,u_x + 0.1\,u_{xx}$ on a periodic domain. Measurements of the field can replace the simulation when they are arranged into the same feature and target arrays.
+
+
+Data generation code
+
+```python
+import numpy as np
+
+L, nx, nu = 2 * np.pi, 128, 0.1
+x = np.linspace(0.0, L, nx, endpoint=False)
+dx = L / nx
+k = 2 * np.pi * np.fft.rfftfreq(nx, d=dx)
+
+def d_dx(u, order=1):
+ return np.real(np.fft.irfft((1j * k) ** order * np.fft.rfft(u), n=nx))
+
+def rhs(u):
+ return -u * d_dx(u) + nu * d_dx(u, order=2)
+
+u = -np.sin(x)
+snapshots = [u.copy()]
+dt = 1e-3
+for i in range(1, 4001):
+ k1 = rhs(u)
+ k2 = rhs(u + 0.5 * dt * k1)
+ k3 = rhs(u + 0.5 * dt * k2)
+ k4 = rhs(u + dt * k3)
+ u = u + (dt / 6.0) * (k1 + 2 * k2 + 2 * k3 + k4)
+ if i % 100 == 0:
+ snapshots.append(u.copy())
+
+U = np.array(snapshots) # shape (41, 128)
+t = np.arange(U.shape[0]) * 0.1
+```
+
+
+
+Estimate spatial derivatives with the smoothing differentiator `savgol_filter` and time derivatives with `np.gradient`. Supply the coordinate spacings and keep each space-time point aligned across the feature columns and target:
+
+```python
+import numpy as np
+
+from scipy.signal import savgol_filter
+
+ux = savgol_filter(U, 21, 3, deriv=1, delta=dx, axis=-1)
+uxx = savgol_filter(U, 21, 3, deriv=2, delta=dx, axis=-1)
+ut = np.gradient(U, t, axis=0)
+
+X = np.stack([U, ux, uxx], axis=-1).reshape(-1, 3)
+y = ut.reshape(-1)
+```
+
+`complexity_of_variables=[1, 2, 3]` makes higher spatial derivatives more expensive, biasing the search toward lower-order terms when candidate fits compete. A successful low-complexity search should place $0.1\,u_{xx} - u\,u_x$ on the Pareto front, matching the right-hand side that generated the snapshots.
+
+```python
+from pysr import PySRRegressor
+
+model = PySRRegressor(
+ binary_operators=["+", "-", "*"],
+ complexity_of_variables=[1, 2, 3],
+ maxsize=20,
+ niterations=100,
+)
+model.fit(X, y, variable_names=["u", "u_x", "u_xx"])
+print(model)
+```
+
+The regression uses estimated derivatives, so simulate the discovered PDE forward from a held-out initial condition before relying on its coefficients.
+
+If the field $u$ has units $\mathrm{m\,s^{-1}}$, then $u_x$ has units $\mathrm{s^{-1}}$, $u_{xx}$ has units $\mathrm{m^{-1}\,s^{-1}}$, and $u_t$ has units $\mathrm{m\,s^{-2}}$. For this feature order, pass the corresponding strings through `X_units=["m/s", "s^-1", "m^-1*s^-1"]` and `y_units="m*s^-2"`.
diff --git a/docs/src/examples/search-behaviour.md b/docs/src/examples/search-behaviour.md
new file mode 100644
index 000000000..e46f0cb7f
--- /dev/null
+++ b/docs/src/examples/search-behaviour.md
@@ -0,0 +1,270 @@
+# Search behaviour
+
+## Automatic batching on a large dataset
+
+`batching` reduces the rows used to score population members during evolution. Hall-of-fame comparisons continue to use every row. With `batching=False`, every evolutionary comparison uses all rows. When `batch_size` is `None`, automatic batching can use every row on small datasets. PySR chooses all rows through 1000, 128 below 5000, 256 below 50,000, and 512 from 50,000 onward. At 20,000 rows, evolution therefore sees 256 rows per comparison, while hall-of-fame scoring sees all 20,000. `batching="auto"` is the default; this example names it to make the choice explicit.
+
+The target depends on two of five input columns. The remaining three inputs act as distractors.
+
+
+Data generation code
+
+```python
+import numpy as np
+
+rng = np.random.default_rng(0)
+X = rng.uniform(-3, 3, (20000, 5))
+y = 2.5382 * np.cos(X[:, 3]) + X[:, 0] ** 2 - 0.5
+```
+
+
+
+```python
+from pysr import PySRRegressor
+
+model = PySRRegressor(
+ binary_operators=["+", "-", "*", "/"],
+ unary_operators=["cos", "exp"],
+ batching="auto",
+ niterations=40,
+ verbosity=0,
+ random_state=0,
+)
+model.fit(X, y, variable_names=["x0", "x1", "x2", "x3", "x4"])
+print(model.equations_[["complexity", "loss", "equation"]].to_string(index=False))
+```
+
+One recovered expression is:
+
+```
+(cos(x3) * 2.5382) + ((x0 * x0) + -0.5)
+```
+
+Run the complete example with `examples/automatic_batching.py`.
+
+## Operators of any arity
+
+
+
+The generic `operators` argument defines an expression grammar by arity. Each integer key gives the number of arguments accepted by the operators listed under it, so one search can mix unary and binary operations with ternary or wider functions. `operators` cannot be combined with `binary_operators` or `unary_operators`.
+
+The target is a four-way minimum,
+
+$$
+y = \min(x + 2,\; 2 - x,\; \sin(3x) + 1.2,\; 0.8).
+$$
+
+A binary `min` would write this reduction as three nested calls. The custom `min4` writes the same mathematical operation in one call, giving the search a shorter structural route to the target.
+
+
+Data generation code
+
+```python
+import numpy as np
+
+DOMAIN = (-4.0, 4.0)
+
+
+def envelope(x):
+ return np.minimum.reduce(
+ [x + 2.0, 2.0 - x, np.sin(3.0 * x) + 1.2, np.full_like(x, 0.8)]
+ )
+
+
+_x = np.sort(np.random.default_rng(0).uniform(*DOMAIN, 400))
+X = _x.reshape(-1, 1).astype(np.float32)
+y = envelope(_x).astype(np.float32)
+
+_unseen = np.linspace(*DOMAIN, 4001)
+X_UNSEEN = _unseen.reshape(-1, 1).astype(np.float32)
+Y_UNSEEN = envelope(_unseen)
+```
+
+
+
+Custom operator definitions are Julia source strings under their arity keys. A SymPy-capable model also needs a matching entry in `extra_sympy_mappings`; PySR rejects a named custom function during setup when its symbolic image is missing. The mappings here provide the `Piecewise` representation for `ifelse` and the `Min` representation for `min4`.
+
+```python
+import sympy
+
+from pysr import PySRRegressor
+
+model = PySRRegressor(
+ operators={
+ 1: ["sin", "cos"],
+ 2: ["+", "-", "*"],
+ 3: ["ifelse(t, a, b) = t > 0 ? a : b"],
+ 4: ["min4(a, b, c, d) = min(a, b, c, d)"],
+ },
+ # A custom operator is rejected at fit time unless it has a SymPy image.
+ extra_sympy_mappings={
+ "ifelse": lambda t, a, b: sympy.Piecewise((a, t > 0), (b, True)),
+ "min4": lambda a, b, c, d: sympy.Min(a, b, c, d),
+ },
+ niterations=500,
+ random_state=0,
+)
+model.fit(X, y, variable_names=["x"])
+print(model.equations_)
+```
+
+One recovered expression is:
+
+```
+min4(0.8, 2.0 - x, 2.0 + x, 1.1999999 + sin(x * 3.0))
+```
+
+The complete example, `examples/any_arity.py`, checks predictions on an unseen grid with `model.predict(X_UNSEEN, index=i)` for each candidate equation.
+
+## Mutations and plugins
+
+PySR accepts mutation and plugin configuration as `PySRRegressor` arguments. `mutations` maps mutation entries to weights, and `plugins` lists plugin entries. Entries merge with shipped defaults by concrete type: a matching type replaces the default entry, while a new type is appended alongside it. This lets one behavior change without reproducing the complete default set.
+
+The following configuration activates Backsolve and names the adaptive mutation-weight plugin.
+
+```python
+from pysr import (
+ AdaptiveMutationWeightsPlugin,
+ BacksolveMutation,
+ PySRRegressor,
+)
+
+model = PySRRegressor(
+ mutations={BacksolveMutation(): 0.1},
+ plugins=[AdaptiveMutationWeightsPlugin()],
+)
+```
+
+`BacksolveMutation` is present at default weight 0.0, so it is available without being sampled. Assigning weight 0.1 makes it eligible for this search. Backsolve inverts the operators above a selected subtree and fits a replacement from expressions already in the population, helping an additive target assemble from separately discovered pieces.
+
+`AdaptiveMutationWeightsPlugin` is part of the default plugin set. It updates mutation weights from observed improvement rates and resets its runtime state for each `fit`. Naming it in `plugins` replaces the existing entry of that type, which is also where `smoothing`, `floor`, and `reward` can be changed. The resolved list keeps simulated annealing and adaptive parsimony alongside adaptive mutation weights. The two settings work together: the mutation weight makes Backsolve eligible, and the adaptive plugin can change how much of the mutation budget each kind receives.
+
+The runnable configuration uses the small target
+
+$$ y = 2.5\cos(3x) + 0.5x^2 - 1. $$
+
+```python
+import numpy as np
+
+x = np.linspace(-3.0, 3.0, 200)
+X = x.reshape(-1, 1)
+y = 2.5 * np.cos(3.0 * x) + 0.5 * x * x - 1.0
+```
+
+```python
+from pysr import AdaptiveMutationWeightsPlugin, BacksolveMutation, PySRRegressor
+
+model = PySRRegressor(
+ binary_operators=["+", "-", "*"],
+ unary_operators=["cos", "exp"],
+ niterations=20,
+ verbosity=0,
+ mutations={BacksolveMutation(): 0.1},
+ plugins=[AdaptiveMutationWeightsPlugin()],
+ random_state=0,
+)
+model.fit(X, y, variable_names=["x"])
+print(model.equations_)
+```
+
+After fitting, inspect `model.julia_options_.mutations` to see the resolved entries and weights. Run the example with `examples/mutations_and_plugins.py`.
+
+## Adaptive mutation weights
+
+
+
+`AdaptiveMutationWeightsPlugin` adjusts mutation weights from their observed success rates. It is enabled by default, and its learned state resets on each `fit`. Pass an instance through `plugins` to change its settings.
+
+- `smoothing` controls the exponential moving average of success rates.
+- `floor` limits the success-rate ratio used to update a multiplier to $[\text{floor}, 1/\text{floor}]$.
+- `reward` selects whether improvement is measured by `"cost"` or `"loss"`.
+
+This example fits a cosine and a quadratic term:
+
+
+Data generation code
+
+```python
+import numpy as np
+
+rng = np.random.default_rng(20260817)
+X = rng.uniform(-3.0, 3.0, size=(200, 5))
+y = 2.5382 * np.cos(X[:, 3]) + X[:, 0] ** 2 - 0.5
+```
+
+
+
+```python
+from pysr import AdaptiveMutationWeightsPlugin, PySRRegressor
+
+model = PySRRegressor(
+ binary_operators=["+", "-", "*", "/"],
+ unary_operators=["cos", "exp"],
+ maxsize=14,
+ niterations=300,
+ populations=1,
+ plugins=[AdaptiveMutationWeightsPlugin(smoothing=0.02, floor=0.05, reward="cost")],
+ random_state=0,
+ verbosity=0,
+)
+model.fit(X, y)
+print(model.equations_)
+```
+
+For an optional diagnostic that records mutation counts and learned multipliers, see `examples/adaptive_mutation_weights.py`.
+
+## The backsolve mutation
+
+
+
+Ordinary mutations change an expression and test the result. Backsolve chooses a non-root subtree and inverts the operators above it to determine the values that subtree would need to produce. It then fits a sparse replacement from expressions already in the population. Greedy forward selection over a basis library finds the replacement within the remaining complexity budget. The replaced subtree can remain in the library, and rows with invalid inverse values are masked. Both `+` and `*` are required for the weighted sum. An event returns a child only when its cost strictly improves on the parent. Each event requires an inversion and a linear solve, so enabling Backsolve can substantially increase search time.
+
+Backsolve applies to ordinary binary scalar expression trees with a real target vector. It rejects unsupported wrappers, shared nodes, root or single-node choices, missing targets, and missing weighted-sum operators. The mutation is experimental and starts at default weight 0.0.
+
+The target combines an oscillation with a quadratic:
+
+$$ y = 2.5\cos(3x) + 0.5x^2 - 1. $$
+
+```python
+import numpy as np
+
+x = np.linspace(-3.0, 3.0, 200)
+X = x.reshape(-1, 1)
+y = 2.5 * np.cos(3.0 * x) + 0.5 * x * x - 1.0
+```
+
+The example omits division and a power operator, so the quadratic must appear as `x * x`. `precision=64` supports the linear solve, and `maxsize=30` leaves room for the complete form.
+
+```python
+from pysr import PySRRegressor
+
+model = PySRRegressor(
+ binary_operators=["+", "-", "*"],
+ unary_operators=["cos", "exp"],
+ weight_backsolve=1.0,
+ maxsize=30,
+ precision=64,
+ niterations=10,
+ verbosity=0,
+ random_state=0,
+)
+model.fit(X, y, variable_names=["x"])
+print(model.equations_[["complexity", "loss", "equation"]].to_string(index=False))
+```
+
+One recovered expression is:
+
+```
+((cos((x + x) + x) * 2.5) + ((x * x) * 0.5)) + -1.0
+```
+
+Check the selected equation against the training target:
+
+```python
+import numpy as np
+
+prediction = np.asarray(model.predict(X), dtype=float)
+print(f"max abs error: {np.max(np.abs(prediction - y)):.3g}")
+```
+
+Use `examples/backsolve_mutation.py` for the complete example.
diff --git a/docs/src/examples/value-types.md b/docs/src/examples/value-types.md
new file mode 100644
index 000000000..53b679410
--- /dev/null
+++ b/docs/src/examples/value-types.md
@@ -0,0 +1,413 @@
+# Value types
+
+## Complex numbers
+
+To fit complex-valued expressions, pass feature and target arrays with a complex NumPy dtype.
+
+The code constructs a complex cosine target and restricts the search to `+`, `-`, `*`, and `cos`:
+
+```python
+import numpy as np
+
+from pysr import PySRRegressor
+
+X = np.random.randn(100, 1) + 1j * np.random.randn(100, 1)
+y = (1 + 2j) * np.cos(X[:, 0] * (0.5 - 0.2j))
+
+model = PySRRegressor(
+ binary_operators=["+", "-", "*"], unary_operators=["cos"], niterations=100,
+)
+
+model.fit(X, y)
+```
+
+The fitted constants belong to the complex-valued expression domain. Use `model.sympy()` to inspect the selected equation in SymPy form:
+
+```python
+model.sympy()
+```
+
+Pass an equation-row index to `predict`; `-1` selects the final row of `equations_`:
+
+```python
+model.predict(X, -1)
+```
+
+## Julia packages and types
+
+PySR delegates its symbolic-regression search to the pure Julia package [SymbolicRegression.jl](https://github.com/astroautomata/SymbolicRegression.jl). A custom operator can call functionality from another Julia package through that backend.
+
+Load the package that supplies the prime lookup before defining the operator. The target in this example is
+
+$$ y = p_{3x + 1} - 5, $$
+
+where $p_i$ denotes the $i$th prime and $x$ is the input feature. The lookup comes from [Primes.jl](https://github.com/JuliaMath/Primes.jl).
+
+First, import the handle to PySR's Julia runtime:
+
+```python
+from pysr import jl
+```
+
+`jl` is the juliacall runtime object, and `jl.seval` evaluates Julia source from Python. Use it to add `Primes.jl` to the active PySR environment:
+
+```python
+from pysr import jl
+
+jl.seval("""
+import Pkg
+Pkg.add("Primes")
+""")
+```
+
+Adding a package and importing its namespace are separate steps. Import `Primes` before defining `p`:
+
+```python
+from pysr import jl
+
+jl.seval("import Primes")
+```
+
+With the package available, define the unary operator used by the search:
+
+```python
+from pysr import jl
+
+jl.seval("""
+function p(i::T) where T
+ if (0.5 < i < 1000)
+ return T(Primes.prime(round(Int, i)))
+ else
+ return T(NaN)
+ end
+end
+""")
+```
+
+The operator rounds inputs within `0.5 < i < 1000` to prime indices and preserves the input type `T`. Outside that domain it returns `T(NaN)` to reject invalid evaluations.
+
+Call the Julia function through `jl.p` to build the prime values used by the synthetic data:
+
+```python
+from pysr import jl
+
+primes = {i: jl.p(i*1.0) for i in range(1, 999)}
+```
+
+Use the lookup table to generate noisy samples of the target:
+
+```python
+import numpy as np
+
+X = np.random.randint(0, 100, 100)[:, None]
+y = [primes[3*X[i, 0] + 1] - 5 + np.random.randn()*0.001 for i in range(100)]
+```
+
+Register `p` with the regular PySR operator interface. The `sympy_p` subclass and `extra_sympy_mappings` give SymPy a symbolic form for the custom function:
+
+```python
+from pysr import PySRRegressor
+import sympy
+
+class sympy_p(sympy.Function):
+ pass
+
+model = PySRRegressor(
+ binary_operators=["+", "-", "*", "/"],
+ unary_operators=["p"],
+ niterations=100,
+ extra_sympy_mappings={"p": sympy_p}
+)
+```
+
+Fit the model on the generated pairs:
+
+```python
+model.fit(X, y)
+```
+
+Since `p` rounds its argument, different real offsets can produce the same integer prime index, so the offset in a fitted expression need not equal `1` exactly. Inspect the selected equation in SymPy with:
+
+```python
+model.sympy()
+```
+
+## Custom value types
+
+`TypeSpec` declares a custom Julia value that can flow through an expression tree. Vectors, tensors, strings, structs, and other Julia value types can serve as payloads when their behavior is described by the type's hooks, operators, and loss.
+
+A model configured with `type_spec=TypeSpec(...)` needs three pieces:
+
+1. Define hooks that sample and change values. For continuously optimized constants, `scalar_constants` exposes scalar coordinates and `with_scalar_constants` rebuilds the value. A custom `mutate` hook handles changes that need a discrete or structural operation.
+2. Choose exactly one of `elementwise_loss`, `loss_function`, or `loss_function_expression`. The selected loss must accept the custom values and return a real-valued score. Elementwise-loss return types are inferred; a full objective requires an explicit `loss_type`.
+3. Supply operators through `operators={...}`, grouped by arity. Each operator must accept and return the declared value type so the type-stability check succeeds. Add an explicit return annotation such as `::Vec2` when Julia cannot infer it.
+
+The examples below combine these pieces for vector values, strings as discrete constants, and payloads with several shapes.
+
+### Vector-valued expression trees
+
+This example searches for an expression whose values are two-dimensional vectors:
+
+$$
+y = \operatorname{rotate90}(x_1) + 2x_2 +
+\begin{bmatrix}0.5 \\ -1.0\end{bmatrix}.
+$$
+
+The 128-by-2 DataFrame has two vector-valued features per row, and each entry of `y` is one target vector. The table dimensions describe samples and features, while each cell of the object array holds a single logical vector value.
+
+```python
+import numpy as np
+import pandas as pd
+
+from pysr import PySRRegressor, TypeSpec
+
+rng = np.random.default_rng(0)
+x1 = [rng.normal(size=2) for _ in range(128)]
+x2 = [rng.normal(size=2) for _ in range(128)]
+X = pd.DataFrame({"x1": x1, "x2": x2})
+y = np.empty(128, dtype=object)
+offset = np.array([0.5, -1.0])
+y[:] = [np.array([-a[1], a[0]]) + 2 * b + offset for a, b in zip(x1, x2)]
+```
+
+Declare `Vec2` with one `Vector{Float64}` field. `sample` creates a random payload of length two, while `scalar_constants` and `with_scalar_constants` expose and rebuild the entries used by continuous constant optimization:
+
+```python
+from pysr import TypeSpec
+
+type_spec = TypeSpec(
+ "Vec2",
+ fields={"data": "Vector{Float64}"},
+ sample="rng -> Vec2(randn(rng, 2))",
+ scalar_constants="value -> value.data",
+ with_scalar_constants=(
+ "(value, scalar_constants) -> Vec2(scalar_constants)"
+ ),
+)
+```
+
+PySR derives how values are initialized, mutated, validated, counted, packed, and unpacked from these hooks. Add `init` or `mutate` for specialized search behavior, or `string` for custom display, when the derived defaults are insufficient.
+
+`TypeSpec` generates the Julia value type from that declaration. Its equivalent Julia struct has one vector field:
+
+```julia
+struct Vec2
+ data::Vector{Float64}
+end
+```
+
+Define operators that consume and return `Vec2` values:
+
+```python
+operators = {
+ 1: [
+ "rotate90(a::Vec2) = Vec2([-a.data[2], a.data[1]])",
+ "double(a::Vec2) = Vec2(2a.data)",
+ ],
+ 2: ["add_vectors(a::Vec2, b::Vec2) = Vec2(a.data + b.data)"],
+}
+```
+
+The elementwise loss compares vector payloads by summing squared differences, which gives the real-valued score the search requires:
+
+```python
+from pysr import PySRRegressor
+
+model = PySRRegressor(
+ type_spec=type_spec,
+ operators=operators,
+ elementwise_loss="vector_loss(a::Vec2, b::Vec2) = sum(abs2, a.data - b.data)",
+ niterations=40,
+ populations=4,
+ maxsize=10,
+)
+
+model.fit(X, y)
+print(model.equations_)
+```
+
+The search space can represent the target as `add_vectors(add_vectors(rotate90(x1), double(x2)), Vec2([0.5, -1.0]))`. The vector constant carries both offset components, which the scalar hook pair exposes for joint optimization.
+
+
+
+String-valued expressions and discrete constants
+
+This example treats each string as a value and chooses a separator from a finite set during mutation. It lowercases the first input, uppercases the second, and joins them with the selected separator:
+
+```python
+import numpy as np
+import pandas as pd
+
+from pysr import PySRRegressor, TypeSpec
+
+X = pd.DataFrame(
+ {
+ "first": ["Py", "symbolic", "hello", "left"],
+ "second": ["SR", "regression", "world", "right"],
+ }
+)
+y = np.array(
+ [f"{a.lower()}-{b.upper()}" for a, b in X.itertuples(index=False)],
+ dtype=object,
+)
+
+type_spec = TypeSpec(
+ "StringValue",
+ fields={"data": "String"},
+ sample='rng -> StringValue(rand(rng, ("", "-", "_")))',
+ mutate="""
+ mutate_string(rng, value, temperature) = StringValue(rand(rng, ("", "-", "_")))
+ """,
+)
+
+model = PySRRegressor(
+ type_spec=type_spec,
+ operators={
+ 1: [
+ "string_lowercase(x::StringValue) = StringValue(lowercase(x.data))",
+ "string_uppercase(x::StringValue) = StringValue(uppercase(x.data))",
+ ],
+ 2: [
+ "string_concat(a::StringValue, b::StringValue) = StringValue(a.data * b.data)"
+ ],
+ },
+ elementwise_loss="""
+ string_loss(a::StringValue, b::StringValue) = Float64(Base.editdistance(a.data, b.data))
+ """,
+ niterations=40,
+)
+
+model.fit(X, y)
+print(model.equations_)
+```
+
+The `sample` and `mutate` hooks choose discrete separator constants. Without a
+`scalar_constants` and `with_scalar_constants` pair, BFGS leaves them unchanged.
+Edit distance provides the real-valued loss for string outputs.
+
+
+
+
+Advanced: recovering a neural network with tensor constants
+
+A `TypeSpec` can place scalar, vector, and matrix payloads in the same Julia value type. Its scalar-constant hooks flatten each payload for BFGS and rebuild its original shape.
+
+The advanced example searches for this two-layer neural network:
+
+$$ y = W_2\operatorname{relu}(W_1x + b_1) + b_2 $$
+
+The targets are vectors. `safe_matmul` and `safe_add` return `NaN` for incompatible shapes, so expressions using those helpers receive an invalid value instead of a dimension error:
+
+```python
+import numpy as np
+import pandas as pd
+
+from pysr import PySRRegressor, TypeSpec
+
+preamble = """
+const NNPayload = Union{Float64, Vector{Float64}, Matrix{Float64}}
+
+safe_matmul(a::Matrix{Float64}, b::Vector{Float64}) =
+ size(a, 2) == length(b) ? a * b : NaN
+safe_matmul(::NNPayload, ::NNPayload) = NaN
+
+safe_add(a::Float64, b::Float64) = a + b
+safe_add(a::T, b::T) where {T<:Union{Vector{Float64}, Matrix{Float64}}} =
+ size(a) == size(b) ? a + b : NaN
+safe_add(::NNPayload, ::NNPayload) = NaN
+
+# Constants sample a random rank, generating only the payload that was chosen:
+function random_nn_payload(rng)
+ rank = rand(rng, 0:2)
+ rank == 0 ? randn(rng) : rank == 1 ? randn(rng, 2) : randn(rng, 2, 2)
+end
+"""
+
+type_spec = TypeSpec(
+ "NNValue",
+ fields={"data": "NNPayload"},
+ sample="rng -> NNValue(random_nn_payload(rng))",
+ # Mutations usually perturb every scalar in the payload, but occasionally
+ # resample a fresh rank:
+ mutate="""
+ (rng, value, temperature) -> if rand(rng) < 0.1
+ NNValue(random_nn_payload(rng))
+ else
+ NNValue(value.data .+ temperature .* randn(rng, size(value.data)...))
+ end
+ """,
+ scalar_constants="""
+ function scalar_constants(value)
+ return value.data isa Float64 ? [value.data] : vec(value.data)
+ end
+ """,
+ with_scalar_constants="""
+ function with_scalar_constants(value, scalar_constants)
+ data = value.data isa Float64 ? scalar_constants[1] :
+ reshape(collect(scalar_constants), size(value.data))
+ return NNValue(data)
+ end
+ """,
+ preamble=preamble,
+)
+```
+
+`preamble` makes the payload union and helpers available before the generated type is
+defined. Mutation can change a constant's rank; continuous optimization preserves its
+shape through the scalar hook pair.
+
+Generate vector-valued targets from fixed weights and biases:
+
+```python
+import numpy as np
+import pandas as pd
+
+rng = np.random.default_rng(0)
+x_values = rng.normal(size=(64, 2))
+W1 = np.array([[1.2, -0.7], [0.5, 1.1]])
+b1 = np.array([0.3, -0.2])
+W2 = np.array([[0.8, -1.0], [1.3, 0.4]])
+b2 = np.array([-0.4, 0.2])
+y_values = (W2 @ np.maximum(x_values @ W1.T + b1, 0).T).T + b2
+
+X = pd.DataFrame({"x": list(x_values)})
+y = pd.Series(list(y_values), dtype=object)
+```
+
+Register the neural-network operations and score matching vector outputs with mean squared error. The loss penalizes other output shapes:
+
+```python
+from pysr import PySRRegressor
+
+model = PySRRegressor(
+ type_spec=type_spec,
+ operators={
+ 1: ["nn_relu(a) = NNValue(max.(a.data, 0.0))"],
+ 2: [
+ "nn_matmul(a, b) = NNValue(safe_matmul(a.data, b.data))",
+ "nn_add(a, b) = NNValue(safe_add(a.data, b.data))",
+ ],
+ },
+ elementwise_loss="""
+ function nn_mse(a, b)::Float64
+ valid = a.data isa Vector && b.data isa Vector && size(a.data) == size(b.data)
+ return valid ? sum(abs2, a.data .- b.data) / length(a.data) : 1.0e6
+ end
+ """,
+ niterations=100,
+ populations=4,
+ maxsize=11,
+)
+
+model.fit(X, y)
+print(model.equations_)
+```
+
+A recovered expression can have the form `nn_matmul(W2, nn_add(b, nn_relu(nn_matmul(W1, nn_add(x, c)))))`. The inner constant `c` represents the first bias through $b_1 = W_1c$, and the outer constant `b` represents the second through $b_2 = W_2b$. Each displayed constant is an `NNValue` containing its fitted scalar, vector, or matrix payload.
+
+
+
+The same custom-value interface can seed a search with `guesses`. In the vector example above,
+`Vec2` is the generated constructor from its `TypeSpec`, `add_vectors` is the declared binary
+operator, and `x1` is the first named feature column. A matching guess is
+`guesses=["add_vectors(x1, Vec2([1.0, 2.0]))"]`.
diff --git a/docs/src/migration.md b/docs/src/migration.md
new file mode 100644
index 000000000..3fa3b7437
--- /dev/null
+++ b/docs/src/migration.md
@@ -0,0 +1,207 @@
+# Migrating from v1 to v2
+
+PySR v2 upgrades the Julia backend from SymbolicRegression.jl v1.11 to v2.
+Most `PySRRegressor` constructor calls still work, but several search defaults
+changed. The old parametric-expression interface and v1 checkpoints do not
+carry over.
+
+Before upgrading, keep a copy of each run directory, including
+`hall_of_fame.csv`. Install v2 in a new Python environment if you need to keep a
+v1 environment available for old pickles or exact search reproduction.
+
+## Search defaults changed
+
+The following changes affect ordinary `PySRRegressor(...)` searches:
+
+| Setting | v1 default | v2 default |
+| --- | --- | --- |
+| `batching` | `False` | `"auto"` |
+| `batch_size` | `50` | `None` |
+| `annealing` | `False` | `True` |
+| `crossover_probability` | `0.0259` | `0.2` |
+| Mutation selection | Fixed v1 weights | New backend weights, adapted during the search |
+
+With `batching="auto"`, the backend enables batching for large datasets while
+still evaluating hall-of-fame candidates on the full dataset. With
+`batch_size=None`, it uses the full dataset for at most 1000 rows, then batches
+of 128, 256, or 512 as the dataset grows.
+
+The default mutation mix changed, including a new feature mutation that can
+replace the input referenced by a variable node. The
+`AdaptiveMutationWeightsPlugin` also updates mutation weights from observed
+success rates. A `weight_*` argument now supplies an initial weight, rather than
+a fixed weight, while that plugin is active.
+
+For a search close to the v1 defaults, set the changed values explicitly and
+omit the adaptive-mutation plugin:
+
+```python
+from pysr import AdaptiveParsimonyPlugin, PySRRegressor
+
+model = PySRRegressor(
+ batching=False,
+ batch_size=50,
+ annealing=False,
+ crossover_probability=0.0259,
+ default_plugins=[AdaptiveParsimonyPlugin()],
+ weight_add_node=2.47,
+ weight_insert_node=0.0112,
+ weight_delete_node=0.87,
+ weight_do_nothing=0.273,
+ weight_mutate_constant=0.0346,
+ weight_mutate_operator=0.293,
+ weight_mutate_feature=0.0,
+ weight_swap_operands=0.198,
+ weight_rotate_tree=4.26,
+ weight_randomize=0.000502,
+ weight_simplify=0.00209,
+ weight_optimize=0.0,
+ weight_backsolve=0.0,
+)
+```
+
+This preserves the v1 search configuration, but it does not make a run
+bit-for-bit identical. Backend implementation changes still alter random-number
+consumption and search trajectories.
+
+Two other backend changes can alter results:
+
+- Simplification now recomputes an expression's cost from scratch. Equivalent
+ expressions can therefore appear in a different hall-of-fame order.
+- Constant optimization can restart zero-valued constants, allowing trajectories
+ that v1 could not reach.
+
+Loss functions are now validated when `fit` is called. An `elementwise_loss` or
+`loss_function` with the wrong number of arguments raises a `ValueError` before
+the search starts.
+
+## Parametric expressions
+
+`ParametricExpressionSpec` and the `category=` argument to `fit` were removed.
+Use `TemplateExpressionSpec.parameters`, append the category to `X`, and index
+parameters from the template.
+
+A v1 model such as:
+
+```python
+from pysr import ParametricExpressionSpec, PySRRegressor
+
+model = PySRRegressor(
+ expression_spec=ParametricExpressionSpec(max_parameters=2),
+)
+model.fit(X, y, category=category)
+```
+
+becomes:
+
+```python
+import numpy as np
+from pysr import PySRRegressor, TemplateExpressionSpec
+
+n_categories = len(np.unique(category))
+spec = TemplateExpressionSpec(
+ combine="f(x1, x2, p1[category], p2[category])",
+ expressions=["f"],
+ variable_names=["x1", "x2", "category"],
+ parameters={"p1": n_categories, "p2": n_categories},
+)
+
+# Julia indexing starts at 1. Add one if Python categories start at 0.
+X_with_category = np.column_stack([X, category + 1])
+model = PySRRegressor(expression_spec=spec)
+model.fit(X_with_category, y)
+```
+
+The category column must also be present when calling `predict`.
+
+## Template expressions
+
+The deprecated `function_symbols=...` and positional constructor forms of
+`TemplateExpressionSpec` were replaced by explicit template fields. State the
+combining expression, inner expressions, and variable names directly:
+
+```python
+spec = TemplateExpressionSpec(
+ combine="sin(f(x1, x2)) + g(x3)^2",
+ expressions=["f", "g"],
+ variable_names=["x1", "x2", "x3"],
+)
+```
+
+This form also supports optimizable parameter vectors through
+`parameters={"p": length}`. See [Template expressions](examples.md#template-expressions)
+for complete examples.
+
+## Saved models and checkpoints
+
+v2 uses checkpoint schema 3. A v1 pickle has no compatible schema marker, so
+`pickle.load(...)` and `PySRRegressor.from_file(...)` reject it with an
+`Unsupported PySR checkpoint schema` error. Warm starts cannot cross the v1 to
+v2 boundary. Re-run the search under v2 if you need to continue evolution.
+
+For standard expressions, you can reconstruct the discovered equations from a
+copy of `hall_of_fame.csv`. Put the CSV in a run directory without the
+incompatible `checkpoint.pkl`, then provide the original operators and feature
+metadata:
+
+```python
+model = PySRRegressor.from_file(
+ run_directory="v1_csv_copy",
+ binary_operators=["+", "-", "*", "/"],
+ unary_operators=["sin", "cos"],
+ n_features_in=3,
+ feature_names_in=["x0", "x1", "x2"],
+)
+```
+
+Custom expression specifications may require a fresh search because their Julia
+objects cannot be reconstructed from the CSV alone.
+
+## v1 spellings that still work
+
+These forms emit a deprecation warning but retain their old behavior:
+
+- camelCase keyword arguments such as `fractionReplaced`, `npop`,
+ `weightMutateConstant`, and `crossoverProbability`; use snake_case instead.
+- `loss=` and `full_objective=`; use `elementwise_loss=` and `loss_function=`.
+- `multithreading=True` or `False`; use `parallelism="multithreading"` or
+ `parallelism="serial"`. Use `parallelism="multiprocessing", procs=N` for
+ process workers.
+- `.equations` and `.raw_julia_state_`; use `.equations_` and `.julia_state_`.
+- The `pysr()` function, `pysr.install()`, and `pysr.init_julia()`; use
+ `PySRRegressor`. Dependencies are installed when PySR is imported.
+
+These constructor arguments are accepted but ignored with a warning:
+
+- `julia_project` and `julia_kwargs`. Configure the backend through `juliapkg`;
+ see [Customizing the backend](backend.md).
+- `weights`, `variable_names`, and `Xresampled`. Pass them to `fit` instead.
+
+## New APIs in v2
+
+- **N-ary operators:** use
+ `operators={1: ["sin"], 2: ["+", "*"], 3: ["clamp"]}`. The legacy
+ `binary_operators` and `unary_operators` arguments still work, but cannot be
+ combined with `operators`. Constraints for n-ary operators use n-tuples, such
+ as `{"clamp": (-1, -1, 1)}`.
+- **First-class mutations, crossovers, and plugins:** configure weighted
+ instances with `mutations=` and `crossovers=`, and compose search behavior
+ with `plugins=`. See the [API reference](api.md#mutations).
+- **Custom value types:** `TypeSpec` supports expressions over strings, vectors,
+ tensors, structs, and other Julia value types. See
+ [Custom value types](examples.md#custom-value-types).
+- **Template parameters:** `TemplateExpressionSpec.parameters` supports learned
+ scalar or category-indexed parameter vectors.
+- **Equation guesses:** seed a search with `guesses=["sin(x1 * 2.1)", ...]`.
+ `fraction_replaced_guesses` controls their injection rate.
+- **Autodiff backends:** `autodiff_backend="Enzyme"` is available as an
+ experimental alternative to `"Zygote"`.
+- **Worker controls:** `worker_imports` loads Julia packages on multiprocessing
+ workers, and `worker_timeout` controls worker restart timeouts.
+- **Additional mutations:** `weight_mutate_feature` controls feature mutation.
+ Backsolve is disabled by default and can be configured with
+ `BacksolveMutation(...)` through `mutations=`.
+
+## Environment changes
+
+- The minimum versions are Python 3.10 and Julia 1.10.
diff --git a/docs/src/options.md b/docs/src/options.md
index 7a8f5c11e..3590808bc 100644
--- a/docs/src/options.md
+++ b/docs/src/options.md
@@ -245,7 +245,7 @@ The default loss is mean-square error, and weighted mean-square error.
One can pass an arbitrary Julia string to define a custom loss, using,
e.g., `elementwise_loss="myloss(x, y) = abs(x - y)^1.5"`. For more details,
see the
-[Losses](https://milescranmer.github.io/SymbolicRegression.jl/dev/losses/)
+[Losses](https://julia.pysr.ai/dev/losses)
page for SymbolicRegression.jl.
Here are some additional examples:
@@ -276,7 +276,7 @@ model = PySRRegressor(..., weights=weights, elementwise_loss="myloss(x, y, w) =
model.fit(..., weights=weights)
```
-Built-in loss (faster) (see [losses](https://ai.damtp.cam.ac.uk/symbolicregression/dev/losses/)).
+Built-in loss (faster) (see [losses](https://julia.pysr.ai/dev/losses)).
This one computes the L3 norm:
```python
diff --git a/docs/src/public/apple-touch-icon.png b/docs/src/public/apple-touch-icon.png
new file mode 100644
index 000000000..439235373
Binary files /dev/null and b/docs/src/public/apple-touch-icon.png differ
diff --git a/docs/src/public/favicon-16x16.png b/docs/src/public/favicon-16x16.png
new file mode 100644
index 000000000..d6a26a31d
Binary files /dev/null and b/docs/src/public/favicon-16x16.png differ
diff --git a/docs/src/public/favicon-32x32.png b/docs/src/public/favicon-32x32.png
new file mode 100644
index 000000000..d46b06535
Binary files /dev/null and b/docs/src/public/favicon-32x32.png differ
diff --git a/docs/src/public/favicon.ico b/docs/src/public/favicon.ico
new file mode 100644
index 000000000..ba1cdba34
Binary files /dev/null and b/docs/src/public/favicon.ico differ
diff --git a/docs/src/public/favicon.png b/docs/src/public/favicon.png
index 5fcc0ab9f..69498a477 100644
Binary files a/docs/src/public/favicon.png and b/docs/src/public/favicon.png differ
diff --git a/docs/src/slurm.md b/docs/src/slurm.md
new file mode 100644
index 000000000..ec1c395bf
--- /dev/null
+++ b/docs/src/slurm.md
@@ -0,0 +1,50 @@
+# Slurm (Multi-node)
+
+PySR supports running across multiple nodes in an existing Slurm allocation with
+`cluster_manager="slurm"`. Request the resources with `sbatch` or `salloc`, then
+PySR will launch Julia workers across the allocation using
+`SlurmClusterManager.jl`.
+
+For example, this job requests three workers on each of two nodes:
+
+```bash
+#!/bin/bash
+#SBATCH --job-name=pysr
+#SBATCH --partition=normal
+#SBATCH --nodes=2
+#SBATCH --ntasks-per-node=3
+#SBATCH --time=01:00:00
+
+set -euo pipefail
+python pysr_script.py
+```
+
+The corresponding `pysr_script.py` is:
+
+```python
+import numpy as np
+from pysr import PySRRegressor
+
+X = np.random.RandomState(0).randn(1000, 2)
+y = X[:, 0] + 2 * X[:, 1]
+
+model = PySRRegressor(
+ niterations=200,
+ populations=2,
+ parallelism="multiprocessing",
+ cluster_manager="slurm",
+ procs=6,
+)
+model.fit(X, y)
+print(model)
+```
+
+Submit the job with:
+
+```bash
+sbatch pysr_job.sh
+```
+
+`procs` must equal the allocation's total task count, given by `--ntasks` or
+`--nodes` multiplied by `--ntasks-per-node`. Run the Python script once inside
+the allocation; do not wrap it in `srun`.
diff --git a/docs/src/tuning.md b/docs/src/tuning.md
index 443789e49..ec995d29b 100644
--- a/docs/src/tuning.md
+++ b/docs/src/tuning.md
@@ -6,11 +6,11 @@ First, my general tips would be to avoid using redundant operators, like how `po
When running PySR, I usually do the following:
-I run from IPython (Jupyter Notebooks don't work as well[^1]) on the head node of a slurm cluster. Passing `cluster_manager="slurm"` will make PySR set up a run over the entire allocation. I set `procs` equal to the total number of cores over my entire allocation.
+I run from IPython (Jupyter Notebooks don't work as well[^1]) on the head node of a slurm cluster. Passing `cluster_manager="slurm"` will make PySR set up a run over the entire allocation. I set `procs` equal to the total number of tasks across my entire allocation (see the [Slurm page](slurm.md) for a complete multi-node example).
-I use the [tensorboard feature](https://ai.damtp.cam.ac.uk/pysr/examples/#12-using-tensorboard-for-logging) for experiment tracking.
+I use the [tensorboard feature](/examples/instrumentation#using-tensorboard-for-logging) for experiment tracking.
-[^1]: Jupyter Notebooks are supported by PySR, but miss out on some useful features available in IPython and Python: the progress bar, and early stopping with "q". In Jupyter you cannot interrupt a search once it has started; you have to restart the kernel. See [this issue](https://github.com/astroautomata/PySR/issues/260) for updates.
+[^1]: Jupyter Notebooks are supported by PySR, but miss out on some useful features available in IPython and Python: the progress bar, and early stopping with "q". Interrupting a search in Jupyter returns the equations found so far.
1. I start by using the default parameters.
2. I use only the operators I think it needs and no more.
diff --git a/docs/test_deploy_versions.jl b/docs/test_deploy_versions.jl
new file mode 100644
index 000000000..886bd939d
--- /dev/null
+++ b/docs/test_deploy_versions.jl
@@ -0,0 +1,163 @@
+using Test
+
+include(joinpath(@__DIR__, "deploy_versions.jl"))
+
+@testset "stable ownership" begin
+ tags = [v"2.1.0", v"2.0.0", v"1.5.10"]
+
+ @test claims_stable("v2.1.0"; all_tagged_versions = tags, log = false)
+ @test !claims_stable("v2.0.0"; all_tagged_versions = tags, log = false)
+ # determine_bases would report this alpha as owning v2.0.0
+ @test !claims_stable("v2.0.0a2"; all_tagged_versions = tags, log = false)
+ @test !claims_stable("dev"; all_tagged_versions = tags, log = false)
+ @test !claims_stable("previews/PR123"; all_tagged_versions = tags, log = false)
+end
+
+@testset "stable deployment root" begin
+ @test stable_deploy_root("/tmp/gh-pages/stable") == "/tmp/gh-pages"
+ @test stable_deploy_root("/tmp/gh-pages/stable/") == "/tmp/gh-pages"
+end
+
+@testset "stable redirects" begin
+ pages = ["index", "examples", "404", "api"]
+ files = stable_redirect_files(pages, "v2.2.1")
+
+ @test files["api.html"] == """
+
+"""
+ @test files["api/index.html"] == """
+
+"""
+ @test files["stable/api.html"] == """
+
+"""
+ @test files["stable/index.html"] == """
+
+"""
+
+ @test sort(collect(keys(files))) == [
+ "api.html",
+ "api/index.html",
+ "examples.html",
+ "examples/index.html",
+ "stable/api.html",
+ "stable/examples.html",
+ "stable/index.html",
+ ]
+ @test collect(files) == collect(stable_redirect_files(reverse(pages), "v2.2.1"))
+ root_stubs = filter(((path, _),) -> !startswith(path, "stable/"), files)
+ @test all(!occursin(r"v\d", content) for content in values(root_stubs))
+end
+
+@testset "stale redirects" begin
+ files = stable_redirect_files(["api"], "v2.2.1")
+ root_stubs = filter(((path, _),) -> !startswith(path, "stable/"), files)
+
+ root = mktempdir()
+ for (path, content) in root_stubs
+ mkpath(joinpath(root, Base.dirname(path)))
+ write(joinpath(root, path), content)
+ end
+ # A page dropped since the previous release, in both of its stub forms.
+ for (path, content) in stable_redirect_files(["gone"], "v2.2.0")
+ startswith(path, "stable/") && continue
+ mkpath(joinpath(root, Base.dirname(path)))
+ write(joinpath(root, path), content)
+ end
+ # Everything a deployment leaves at the root besides our stubs.
+ write(joinpath(root, "index.html"), "\n")
+ write(joinpath(root, "versions.js"), "var DOC_VERSIONS = [];\n")
+ write(joinpath(root, "CNAME"), "pysr.ai\n")
+ for version in ("v2.2.1", "dev", "stable")
+ mkpath(joinpath(root, version))
+ write(joinpath(root, version, "index.html"), "\n")
+ end
+ write(joinpath(root, "stable", "index.html"), files["stable/index.html"])
+
+ @test stale_root_stubs(root, keys(root_stubs)) == ["gone.html", "gone/index.html"]
+ @test read(joinpath(root, "CNAME"), String) == "pysr.ai\n"
+ @test isempty(stale_root_stubs(root, keys(stable_redirect_files(["api", "gone"], "v2.2.1"))))
+end
+
+@testset "nested redirects" begin
+ files = stable_redirect_files(["api", "examples/objectives", "guide/index"], "v2.3.0")
+
+ @test files["examples/objectives.html"] == """$(STUB_MARKER)
+
+"""
+ @test files["examples/objectives/index.html"] == """$(STUB_MARKER)
+
+"""
+ @test files["stable/examples/objectives.html"] == """$(STUB_MARKER)
+
+"""
+ # A page that is already a directory index gets no `/index.html` of its own.
+ @test files["guide/index.html"] == """$(STUB_MARKER)
+
+"""
+ @test !haskey(files, "guide/index/index.html")
+ @test sort(collect(keys(files))) == [
+ "api.html",
+ "api/index.html",
+ "examples/objectives.html",
+ "examples/objectives/index.html",
+ "guide/index.html",
+ "stable/api.html",
+ "stable/examples/objectives.html",
+ "stable/guide/index.html",
+ "stable/index.html",
+ ]
+ root_stubs = filter(((path, _),) -> !startswith(path, "stable/"), files)
+ @test all(!occursin(r"v\d", content) for content in values(root_stubs))
+ # A flat page and a directory index cannot both own the same stub path.
+ @test_throws ErrorException stable_redirect_files(["guide", "guide/index"], "v2.3.0")
+end
+
+@testset "built pages" begin
+ dist = mktempdir()
+ mkpath(joinpath(dist, "examples"))
+ mkpath(joinpath(dist, "assets"))
+ for path in ("index.html", "404.html", "api.html", "examples/objectives.html")
+ write(joinpath(dist, path), "\n")
+ end
+ write(joinpath(dist, "assets", "app.js"), "")
+ write(joinpath(dist, "hashmap.json"), "{}")
+
+ @test built_pages(dist) == ["404", "api", "examples/objectives", "index"]
+ @test sort(collect(keys(stable_redirect_files(built_pages(dist), "v2.3.0")))) == [
+ "api.html",
+ "api/index.html",
+ "examples/objectives.html",
+ "examples/objectives/index.html",
+ "stable/api.html",
+ "stable/examples/objectives.html",
+ "stable/index.html",
+ ]
+end
+
+@testset "nested stale redirects" begin
+ current = stable_redirect_files(["api"], "v2.3.0")
+ root_stubs = filter(((path, _),) -> !startswith(path, "stable/"), current)
+
+ root = mktempdir()
+ for (path, content) in root_stubs
+ mkpath(joinpath(root, Base.dirname(path)))
+ write(joinpath(root, path), content)
+ end
+ # Stubs for a nested page that has since been removed or renamed.
+ for (path, content) in stable_redirect_files(["examples/objectives"], "v2.2.1")
+ startswith(path, "stable/") && continue
+ mkpath(joinpath(root, Base.dirname(path)))
+ write(joinpath(root, path), content)
+ end
+ # Real builds, whose pages carry no marker and so are never candidates.
+ for version in ("v2.2.1", "dev", "previews/PR1", "stable")
+ mkpath(joinpath(root, version, "examples"))
+ write(joinpath(root, version, "examples", "objectives.html"), "\n")
+ end
+ write(joinpath(root, "previews", "PR1", "stale.html"), STUB_MARKER * "\n")
+ write(joinpath(root, "stable", "stale.html"), STUB_MARKER * "\n")
+
+ @test stale_root_stubs(root, keys(root_stubs)) ==
+ ["examples/objectives.html", "examples/objectives/index.html"]
+end
diff --git a/environment.yml b/environment.yml
index 6ef91fbe5..587f9b410 100644
--- a/environment.yml
+++ b/environment.yml
@@ -2,11 +2,11 @@ name: test
channels:
- conda-forge
dependencies:
- - python>=3.9
- - sympy>=1.0.0,<2.0.0
+ - python>=3.10
+ - sympy>=1.14.0,<2.0.0
- pandas>=0.21.0,<4.0.0
- numpy>=1.13.0,<3.0.0
- scikit-learn>=1.0.0,<2.0.0
- - pyjuliacall>=0.9.28,<0.9.29
+ - pyjuliacall>=0.9.29,<0.9.36
- click>=7.0.0,<9.0.0
- typing-extensions>=4.0.0,<5.0.0
diff --git a/examples/adaptive_mutation_weights.py b/examples/adaptive_mutation_weights.py
new file mode 100644
index 000000000..2fd92828a
--- /dev/null
+++ b/examples/adaptive_mutation_weights.py
@@ -0,0 +1,305 @@
+"""Watch the search re-allocate its own mutation budget.
+
+`AdaptiveMutationWeightsPlugin` is on by default: per mutation kind it counts attempts
+and strictly improving children, and scales that kind's weight by the learned ratio.
+SymbolicRegression.jl logs none of those counters, so this example puts the shipped
+plugin inside a plugin of its own that forwards every hook to it and reads its state
+on the way out, then prints where the budget ended up.
+"""
+
+from dataclasses import dataclass, field
+
+import numpy as np
+
+from pysr import PySRRegressor
+from pysr.julia_import import AnyValue, jl
+from pysr.plugins import (
+ AbstractPlugin,
+ AdaptiveMutationWeightsPlugin,
+ AdaptiveParsimonyPlugin,
+ SimulatedAnnealingPlugin,
+)
+
+_rng = np.random.default_rng(20260817)
+X = _rng.uniform(-3.0, 3.0, size=(200, 5))
+y = 2.5382 * np.cos(X[:, 3]) + X[:, 0] ** 2 - 0.5
+VARIABLE_NAMES = None
+
+# `DocsTracedAdaptive` holds a real `AdaptiveMutationWeightsPlugin` and forwards
+# `init_plugin_state`, `on_mutation_end!` and `condition_mutation_weights!` to it
+# unchanged, so the adaptation arithmetic during the run is the library's own. The
+# wrapper only counts draws per mutation kind and snapshots the plugin's own
+# `multipliers` vector as it moves.
+_TRACER_JL = """
+if !isdefined(Main, :DocsTracedAdaptive)
+
+const DocsAMW = SymbolicRegression.AdaptiveMutationWeightsModule
+
+struct DocsTracedAdaptive <: SymbolicRegression.AbstractPlugin
+ inner::DocsAMW.AdaptiveMutationWeightsPlugin
+ stride::Int
+end
+
+DocsTracedAdaptive(; stride::Int=1000, kws...) =
+ DocsTracedAdaptive(DocsAMW.AdaptiveMutationWeightsPlugin(; kws...), stride)
+
+mutable struct DocsTracedState
+ inner::DocsAMW.AdaptiveMutationWeightsState
+ draws::Vector{Float64}
+ seen::Int
+ snapshots::Vector{Tuple{Int,Vector{Float64}}}
+ registered::Bool
+end
+
+const DOCS_TRACED_STATES = DocsTracedState[]
+
+function SymbolicRegression.init_plugin_state(p::DocsTracedAdaptive, options, dataset)
+ inner = SymbolicRegression.init_plugin_state(p.inner, options, dataset)
+ n = length(options.mutations)
+ return DocsTracedState(inner, zeros(n), 0, Tuple{Int,Vector{Float64}}[], false)
+end
+
+function SymbolicRegression.on_mutation_end!(
+ s::DocsTracedState,
+ p::DocsTracedAdaptive,
+ mutation::SymbolicRegression.AbstractMutation,
+ event::SymbolicRegression.MutationEvent,
+ dataset,
+ options::SymbolicRegression.AbstractOptions,
+)
+ if !s.registered
+ push!(DOCS_TRACED_STATES, s)
+ s.registered = true
+ end
+ s.draws[event.mutation_idx] += 1.0
+ s.seen += 1
+ SymbolicRegression.on_mutation_end!(s.inner, p.inner, mutation, event, dataset, options)
+ if s.seen % p.stride == 0
+ push!(s.snapshots, (s.seen, copy(s.inner.multipliers)))
+ end
+ return nothing
+end
+
+function SymbolicRegression.condition_mutation_weights!(
+ weights::AbstractVector,
+ s::DocsTracedState,
+ p::DocsTracedAdaptive,
+ member,
+ options::SymbolicRegression.AbstractOptions,
+ curmaxsize,
+ nfeatures,
+)
+ return SymbolicRegression.condition_mutation_weights!(
+ weights, s.inner, p.inner, member, options, curmaxsize, nfeatures
+ )
+end
+
+end
+"""
+
+
+@dataclass(frozen=True)
+class TracedAdaptiveMutationWeights(AbstractPlugin):
+ """`AdaptiveMutationWeightsPlugin`, wrapped so its learned state can be read."""
+
+ stride: int = 1000
+ inner: AdaptiveMutationWeightsPlugin = field(
+ default_factory=AdaptiveMutationWeightsPlugin
+ )
+
+ def julia_plugin(self) -> AnyValue:
+ jl.seval(_TRACER_JL)
+ # PySR calls this once per `fit`, while building that fit's options.
+ jl.seval("empty!(DOCS_TRACED_STATES)")
+ return jl.DocsTracedAdaptive(
+ stride=self.stride,
+ smoothing=self.inner.smoothing,
+ floor=self.inner.floor,
+ reward=jl.Symbol(self.inner.reward),
+ )
+
+
+TRACER = TracedAdaptiveMutationWeights()
+
+MODEL_KWARGS = dict(
+ binary_operators=["+", "-", "*", "/"],
+ unary_operators=["cos", "exp"],
+ maxsize=14,
+ niterations=300,
+ # The plugin keeps one set of counters per population, so a single population is
+ # what makes the printed numbers a run rather than an average over runs.
+ populations=1,
+ # The adaptive plugin is a *default* plugin, so it has to be replaced rather than
+ # added: an extra plugin of a new type leaves the shipped one in place and the
+ # search adapts twice. `alpha=3.17` is PySR's own default, asserted in `main`.
+ plugins=[],
+ default_plugins=[
+ SimulatedAnnealingPlugin(alpha=3.17),
+ AdaptiveParsimonyPlugin(),
+ TRACER,
+ ],
+ verbosity=0,
+)
+
+
+def check(model):
+ """The target recovered to numerical precision."""
+ return bool(model.equations_["loss"].min() < 1e-12)
+
+
+def _states():
+ return int(jl.seval("length(DOCS_TRACED_STATES)"))
+
+
+def _pooled(expr, n):
+ """Sum a per-state vector over every population's state."""
+ total = [0.0] * n
+ for s in range(1, _states() + 1):
+ for i, v in enumerate(jl.seval(f"DOCS_TRACED_STATES[{s}].{expr}")):
+ total[i] += float(v)
+ return total
+
+
+def mutation_statistics(model):
+ """Per-kind counters read out of the plugin's own state after a fit."""
+ jl.DOCS_TRACED_OPTIONS = model.julia_options_
+ names = [
+ str(v)
+ for v in jl.seval(
+ "[string(typeof(first(p))) for p in DOCS_TRACED_OPTIONS.mutations]"
+ )
+ ]
+ base = [
+ float(v)
+ for v in jl.seval("[Float64(last(p)) for p in DOCS_TRACED_OPTIONS.mutations]")
+ ]
+ active = [bool(v) for v in jl.seval("DOCS_TRACED_STATES[1].inner.active")]
+ n = len(names)
+ n_states = _states()
+
+ draws = _pooled("draws", n)
+ attempts = _pooled("inner.attempts", n)
+ successes = _pooled("inner.successes", n)
+ multipliers = [m / n_states for m in _pooled("inner.multipliers", n)]
+ strength = float(
+ jl.seval("DocsAMW.AdaptiveMutationWeightsPlugin().adaptation_strength")
+ )
+
+ # The kinds the run actually drew from: nonzero base weight, and legal for a plain
+ # expression. Effective weight is base * multiplier ** adaptation_strength.
+ drawn = [i for i in range(n) if draws[i] > 0]
+ eff = {i: base[i] * multipliers[i] ** strength for i in drawn}
+ base_total = sum(base[i] for i in drawn)
+ eff_total = sum(eff.values())
+ draw_total = sum(draws[i] for i in drawn)
+
+ rows = [
+ {
+ "mutation": names[i],
+ "index": i + 1,
+ "adapts": active[i],
+ "base_weight": base[i],
+ "draws": int(draws[i]),
+ "attempts": int(attempts[i]),
+ "successes": int(successes[i]),
+ # The quantity the plugin scores with: a Laplace-smoothed improvement rate.
+ "rate": (successes[i] + 1.0) / (attempts[i] + 2.0) if active[i] else None,
+ "multiplier": multipliers[i],
+ "share_shipped": base[i] / base_total,
+ "share_learned": eff[i] / eff_total,
+ "share_drawn": draws[i] / draw_total,
+ }
+ for i in drawn
+ ]
+ rows.sort(key=lambda r: -r["multiplier"])
+ meta = {
+ "n_states": n_states,
+ "n_mutations": int(sum(draws)),
+ "adaptation_strength": strength,
+ "plugins": [
+ str(v)
+ for v in jl.seval(
+ "[string(typeof(p)) for p in DOCS_TRACED_OPTIONS.plugins]"
+ )
+ ],
+ }
+ return rows, meta
+
+
+def multiplier_trajectory(rows, n_points=6):
+ """The learned multipliers at geometrically spaced points through the run.
+
+ Geometric rather than even: most of the re-allocation happens in the first few
+ thousand mutations, which even spacing would step straight over.
+ """
+ n_states = _states()
+ n_snaps = min(
+ int(jl.seval(f"length(DOCS_TRACED_STATES[{s}].snapshots)"))
+ for s in range(1, n_states + 1)
+ )
+ if n_snaps == 0:
+ return []
+ picks = sorted(
+ {1} | {max(1, round(n_snaps ** (k / n_points))) for k in range(1, n_points + 1)}
+ )
+ out = []
+ for p in picks:
+ seen = 0
+ totals = None
+ for s in range(1, n_states + 1):
+ snap = f"DOCS_TRACED_STATES[{s}].snapshots[{p}]"
+ seen += int(jl.seval(f"{snap}[1]"))
+ values = [float(v) for v in jl.seval(f"{snap}[2]")]
+ totals = (
+ values if totals is None else [a + b for a, b in zip(totals, values)]
+ )
+ out.append((seen, [totals[r["index"] - 1] / n_states for r in rows]))
+ return out
+
+
+def report(rows, meta):
+ print(
+ f"\n{meta['n_mutations']} mutations, {meta['n_states']} population state(s), adaptation_strength={meta['adaptation_strength']}, plugins={meta['plugins']}"
+ )
+
+ print(
+ f"\n{'mutation':<23}{'draws':>8}{'improves':>10}{'multiplier':>12}{'share: shipped':>16}{'learned':>10}{'drawn':>9}"
+ )
+ for r in rows:
+ rate = " -" if r["rate"] is None else f"{r['rate']:10.4f}"
+ print(
+ f"{r['mutation']:<23}{r['draws']:>8}{rate}{r['multiplier']:>12.3f}{r['share_shipped']:>16.3%}{r['share_learned']:>10.3%}{r['share_drawn']:>9.3%}"
+ )
+ print("\nshipped and learned are weight shares; drawn is what the search sampled,")
+ print("which also reflects the engine zeroing kinds illegal for a given member.")
+
+ trajectory = multiplier_trajectory(rows)
+ if trajectory:
+ print("\nlearned multiplier through the run, by mutations elapsed:")
+ print(f"{'mutation':<23}" + "".join(f"{seen:>10}" for seen, _ in trajectory))
+ for j, r in enumerate(rows):
+ print(
+ f"{r['mutation']:<23}"
+ + "".join(f"{values[j]:>10.3f}" for _, values in trajectory)
+ )
+
+
+def main():
+ assert MODEL_KWARGS["default_plugins"][0].alpha == PySRRegressor().alpha
+
+ model = PySRRegressor(**MODEL_KWARGS, random_state=0)
+ model.fit(X, y, variable_names=VARIABLE_NAMES)
+ print(model.equations_[["complexity", "loss", "equation"]].to_string(index=False))
+
+ rows, meta = mutation_statistics(model)
+ assert meta["plugins"] == [
+ "SimulatedAnnealingPlugin",
+ "AdaptiveParsimonyPlugin",
+ "DocsTracedAdaptive",
+ ], meta["plugins"]
+ report(rows, meta)
+ print(f"\nrecovered exactly: {check(model)}")
+
+
+if __name__ == "__main__":
+ main()
diff --git a/examples/affine_cipher.py b/examples/affine_cipher.py
new file mode 100644
index 000000000..e17339e3d
--- /dev/null
+++ b/examples/affine_cipher.py
@@ -0,0 +1,69 @@
+"""Break an affine cipher with symbolic regression over a custom letter type.
+
+The cipher is E(p) = (5p + 8) mod 26 on the 26-letter alphabet. The search sees
+only pairs of (ciphertext letter, plaintext letter) and has to discover the
+decryption law. Letters are a custom type rather than floats, so the operators are
+modular letter arithmetic and the constants in the printed expressions are letters.
+"""
+
+import numpy as np
+
+from pysr import PySRRegressor, TypeSpec
+
+ALPHABET = "ABCDEFGHIJKLMNOPQRSTUVWXYZ"
+PLAIN = "THELAWSOFPHYSICSAREWRITTENINTHELANGUAGEOFMATHEMATICS"
+CIPHER = "".join(ALPHABET[(5 * ALPHABET.index(c) + 8) % 26] for c in PLAIN)
+
+X = np.array([[ALPHABET.index(c)] for c in CIPHER], dtype=object)
+y = np.array([ALPHABET.index(c) for c in PLAIN], dtype=object)
+
+# A letter is an integer code in 0:25. Mutation walks one letter up or down, and
+# the string form prints the letter itself, so a constant reads as `V`, not 21.
+SPEC = TypeSpec(
+ "Letter",
+ fields={"code": "Int"},
+ sample="rng -> Letter(rand(rng, 0:25))",
+ mutate="(rng, value, temperature) -> Letter(mod(value.code + rand(rng, (-1, 1)), 26))",
+ string="value -> string(Char('A' + value.code))",
+)
+
+# Distance in letter codes between the prediction and the target, which PySR averages
+# over the rows. It is zero exactly when all 52 letters are right, and a hit-or-miss
+# count would score a candidate one letter off the same as one twenty off.
+LETTER_LOSS = "letter_loss(prediction::Letter, target::Letter)::Float64 = abs(prediction.code - target.code)"
+
+MODEL_KWARGS = dict(
+ type_spec=SPEC,
+ # Modular arithmetic on letters: every operator stays inside the alphabet.
+ operators={
+ 2: [
+ "shift(a::Letter, b::Letter) = Letter(mod(a.code + b.code, 26))",
+ "unshift(a::Letter, b::Letter) = Letter(mod(a.code - b.code, 26))",
+ "mix(a::Letter, b::Letter) = Letter(mod(a.code * b.code, 26))",
+ ]
+ },
+ elementwise_loss=LETTER_LOSS,
+ niterations=20,
+)
+
+
+def decode(model):
+ return "".join(ALPHABET[int(v) % 26] for v in model.predict(X))
+
+
+def check(model):
+ """The whole message comes back, letter for letter."""
+ return decode(model) == PLAIN
+
+
+def main():
+ model = PySRRegressor(**MODEL_KWARGS, random_state=0)
+ model.fit(X, y)
+ print(model.equations_[["complexity", "loss", "equation"]].to_string(index=False))
+ print(f"\ndecoded: {decode(model)}")
+ print(f"plaintext: {PLAIN}")
+ print(f"exact: {check(model)}")
+
+
+if __name__ == "__main__":
+ main()
diff --git a/examples/affine_freemod.py b/examples/affine_freemod.py
new file mode 100644
index 000000000..1e9d052d2
--- /dev/null
+++ b/examples/affine_freemod.py
@@ -0,0 +1,90 @@
+"""Affine cipher with the modulus withheld: 26 has to be learned, not declared.
+
+The data is 52 (ciphertext, plaintext) letter-code pairs from E(p) = (5p + 8) mod 26.
+The search sees only the codes. The value type is a plain integer in 0:51 and the three
+operators are integer `+`, `*` and `mod` with nothing folded in, so the decryption law
+`mod(21*x0 + 14, 26)` has to be written out of three learned constants, the 26 among
+them.
+"""
+
+import numpy as np
+
+from pysr import PySRRegressor, TypeSpec
+
+ALPHABET = "ABCDEFGHIJKLMNOPQRSTUVWXYZ"
+A, B = 5, 8
+PLAIN = "THELAWSOFPHYSICSAREWRITTENINTHELANGUAGEOFMATHEMATICS"
+CIPHER = "".join(ALPHABET[(A * ALPHABET.index(c) + B) % 26] for c in PLAIN)
+
+X = np.array([[ALPHABET.index(c)] for c in CIPHER], dtype=object)
+y = np.array([ALPHABET.index(c) for c in PLAIN], dtype=object)
+
+# The sampled constant range. 26 must not be the only value that fits in it, so the
+# range runs to twice the alphabet and every modulus up to 51 is reachable.
+CMAX = 51
+
+# Constant-proposal jump rate: the share of proposals that resample the whole range
+# instead of stepping by one. One of the three constants is a modulus, and 25 or 27 in
+# that slot scrambles the output as badly as any other wrong value, so a walk alone
+# cannot find it.
+JUMP = 0.35
+
+SPEC = TypeSpec(
+ "Num",
+ fields={"n": "Int"},
+ sample=f"rng -> Num(rand(rng, 0:{CMAX}))",
+ mutate=f"(rng, value, temperature) -> Num(rand(rng) < {JUMP} ? rand(rng, 0:{CMAX}) : clamp(value.n + rand(rng, (-1, 1)), 0, {CMAX}))",
+ string="value -> string(value.n)",
+)
+
+# Integer arithmetic, no alphabet size anywhere. `mod` needs a total extension because
+# a subtree can evaluate to zero; a zero modulus returns the dividend, which keeps the
+# operator total without making any candidate look better than it is.
+OPERATORS = {
+ 2: [
+ "add(a::Num, b::Num) = Num(a.n + b.n)",
+ "mul(a::Num, b::Num) = Num(a.n * b.n)",
+ "modulo(a::Num, b::Num) = Num(b.n == 0 ? a.n : mod(a.n, b.n))",
+ ]
+}
+
+# Absolute distance between the candidate's number and the target letter code, which
+# PySR averages over the rows. It is zero exactly when every one of the 52 letters is
+# right, and a hit-or-miss score would give a candidate that is one off the same loss as
+# one that is twenty off, leaving the constant search nothing to climb.
+NUM_LOSS = (
+ "num_loss(prediction::Num, target::Num)::Float64 = abs(prediction.n - target.n)"
+)
+
+MODEL_KWARGS = dict(
+ type_spec=SPEC,
+ operators=OPERATORS,
+ elementwise_loss=NUM_LOSS,
+ niterations=1000,
+ maxsize=12,
+ should_optimize_constants=False,
+ verbosity=0,
+)
+
+
+def decode(model):
+ return "".join(ALPHABET[int(v) % 26] for v in model.predict(X))
+
+
+def check(model):
+ """The whole message comes back, letter for letter."""
+ return decode(model) == PLAIN
+
+
+def main():
+ model = PySRRegressor(**MODEL_KWARGS, random_state=0)
+ model.fit(X, y)
+
+ print(model.equations_[["complexity", "loss", "equation"]].to_string(index=False))
+ print(f"\ndecoded: {decode(model)}")
+ print(f"plaintext: {PLAIN}")
+ print(f"exact: {check(model)}")
+
+
+if __name__ == "__main__":
+ main()
diff --git a/examples/agent_loop_guesses.py b/examples/agent_loop_guesses.py
new file mode 100644
index 000000000..4ca34f51c
--- /dev/null
+++ b/examples/agent_loop_guesses.py
@@ -0,0 +1,91 @@
+"""Close an agent loop around PySR: read the front, write the next equation, search again.
+
+The data is 160 measurements of kinetic energy against mass and velocity in natural units,
+at velocities high enough that the Newtonian form is only an approximation. A cold search
+climbs a ladder of polynomials in velocity and stops three to four orders of magnitude
+above the measurement noise. Reading the front says what is missing: the relativistic
+law, passed back in through `guesses=`.
+"""
+
+import numpy as np
+
+from pysr import PySRRegressor
+
+rng = np.random.default_rng(42)
+mass = rng.uniform(0.5, 3.0, 160)
+velocity = rng.uniform(0.08, 0.88, 160)
+kinetic_energy_exact = mass * (1 / np.sqrt(1 - velocity**2) - 1)
+NOISE_SIGMA = 0.0015 * float(np.median(kinetic_energy_exact))
+kinetic_energy = kinetic_energy_exact + rng.normal(0.0, NOISE_SIGMA, mass.size)
+NOISE_FLOOR = NOISE_SIGMA**2
+
+X = np.stack([mass, velocity], axis=1)
+y = kinetic_energy
+VARIABLE_NAMES = ["mass", "velocity"]
+
+MODEL_KWARGS = dict(
+ operators={1: ["sqrt"], 2: ["+", "-", "*", "/"]},
+ niterations=10,
+ verbosity=0,
+)
+
+# The filmed session handed the front to a coding agent, which replied with this. Each
+# entry is one round's proposal, so the loop stops when the list runs out.
+PROPOSALS = ["(mass / sqrt(1 - (velocity * velocity))) - mass"]
+
+
+def agent(front, tried):
+ """Stand in for the model call: show the front, return the next expression, or None.
+
+ A real loop prompts a model right here with the front printed below and the list of
+ expressions already tried, and parses one expression out of the reply.
+ """
+ print("front handed to the agent:")
+ print(
+ front[["complexity", "loss", "equation"]].to_string(index=False),
+ "\n",
+ flush=True,
+ )
+ return PROPOSALS[len(tried)] if len(tried) < len(PROPOSALS) else None
+
+
+def check(model):
+ """The search reached the measurement noise floor."""
+ return bool(model.equations_["loss"].min() <= 2 * NOISE_FLOOR)
+
+
+def main(seed=0):
+ guesses, tried, rounds = None, [], []
+ while True:
+ model = PySRRegressor(**MODEL_KWARGS, guesses=guesses, random_state=seed)
+ model.fit(X, y, variable_names=VARIABLE_NAMES)
+ front = model.equations_
+ best = front.loc[front["loss"].idxmin()]
+ print(
+ f"round {len(rounds)} guess: {guesses[0] if guesses else 'none (cold start)'}"
+ )
+ print(f" loss {best['loss']:.4g} at complexity {best['complexity']}")
+ print(f" {best['equation']}\n", flush=True)
+ rounds.append(
+ dict(
+ guess=guesses[0] if guesses else None,
+ loss=float(best["loss"]),
+ complexity=int(best["complexity"]),
+ equation=str(best["equation"]),
+ )
+ )
+ if check(model):
+ break
+ proposal = agent(front, tried)
+ if proposal is None:
+ break
+ tried.append(proposal)
+ guesses = [proposal] # single output: a plain list of expression strings
+
+ print(front[["complexity", "loss", "equation"]].to_string(index=False))
+ print(f"\nnoise floor {NOISE_FLOOR:.4g}; reached: {check(model)}")
+ return rounds
+
+
+if __name__ == "__main__":
+ main()
diff --git a/examples/any_arity.py b/examples/any_arity.py
new file mode 100644
index 000000000..af287ae59
--- /dev/null
+++ b/examples/any_arity.py
@@ -0,0 +1,101 @@
+"""Rediscover a clipped envelope written with a four-argument minimum.
+
+Operators are keyed by arity, so an operator of any arity can be supplied: here a
+three-argument conditional and a four-argument minimum sit side by side in the
+same set. The target `min4(x + 2, 2 - x, sin(3x) + 1.2, 0.8)` folds a rising
+ramp, a falling ramp, a shifted sinusoid and a flat cap together in one call,
+which binary `min` would need three nested calls to write.
+
+Recovery is judged on 4001 points the search never saw, not on the 400 it fitted.
+"""
+
+import numpy as np
+import sympy
+
+from pysr import PySRRegressor
+
+DOMAIN = (-4.0, 4.0)
+
+
+def envelope(x):
+ return np.minimum.reduce(
+ [x + 2.0, 2.0 - x, np.sin(3.0 * x) + 1.2, np.full_like(x, 0.8)]
+ )
+
+
+_x = np.sort(np.random.default_rng(0).uniform(*DOMAIN, 400))
+X = _x.reshape(-1, 1).astype(np.float32)
+y = envelope(_x).astype(np.float32)
+VARIABLE_NAMES = ["x"]
+
+_unseen = np.linspace(*DOMAIN, 4001)
+X_UNSEEN = _unseen.reshape(-1, 1).astype(np.float32)
+Y_UNSEEN = envelope(_unseen)
+
+MODEL_KWARGS = dict(
+ operators={
+ 1: ["sin", "cos"],
+ 2: ["+", "-", "*"],
+ 3: ["ifelse(t, a, b) = t > 0 ? a : b"],
+ 4: ["min4(a, b, c, d) = min(a, b, c, d)"],
+ },
+ # A custom operator is rejected at fit time unless it has a SymPy image.
+ extra_sympy_mappings={
+ "ifelse": lambda t, a, b: sympy.Piecewise((a, t > 0), (b, True)),
+ "min4": lambda a, b, c, d: sympy.Min(a, b, c, d),
+ },
+ niterations=500,
+ verbosity=0,
+)
+
+TOLERANCE = 1e-4
+
+
+def unseen_errors(model):
+ """Largest error of each front member on the points the search never saw."""
+ return {
+ i: float(
+ np.abs(
+ np.asarray(model.predict(X_UNSEEN, index=i), np.float64) - Y_UNSEEN
+ ).max()
+ )
+ for i in model.equations_.index
+ }
+
+
+def check(model):
+ """A front member reproduces the law off the fitted samples, not just on them."""
+ return bool(any(e < TOLERANCE for e in unseen_errors(model).values()))
+
+
+def main():
+ model = PySRRegressor(**MODEL_KWARGS, random_state=0)
+ model.fit(X, y, variable_names=VARIABLE_NAMES)
+
+ front = model.equations_[["complexity", "loss", "equation"]]
+ errors = unseen_errors(model)
+ print(
+ front.assign(unseen_error=[errors[i] for i in front.index]).to_string(
+ index=False
+ )
+ )
+
+ uses = front["equation"].str.contains
+ print(
+ f"\nfront members using the arity-3 ifelse: {int(uses('ifelse').sum())}, using the arity-4 min4: {int(uses('min4').sum())}"
+ )
+
+ print("true law: min4(x + 2, 2 - x, sin(3x) + 1.2, 0.8)")
+ recovered = [i for i, e in errors.items() if e < TOLERANCE]
+ if recovered:
+ best = min(recovered, key=lambda i: front.complexity[i])
+ print(
+ f"recovered at complexity {int(front.complexity[best])}, max error {errors[best]:.3g} over 4001 unseen points:"
+ )
+ print(f" {front.equation[best]}")
+ else:
+ print("no front member agrees with the law off the fitted samples")
+
+
+if __name__ == "__main__":
+ main()
diff --git a/examples/automatic_batching.py b/examples/automatic_batching.py
new file mode 100644
index 000000000..30f7afa6b
--- /dev/null
+++ b/examples/automatic_batching.py
@@ -0,0 +1,46 @@
+"""Recover y = 2.5382*cos(x3) + x0^2 - 0.5 from 20,000 rows.
+
+PySR defaults to `batching="auto"`, which evolves populations against a mini-batch
+while still ranking the hall of fame on the full dataset. With `batch_size=None` the
+batch is chosen from the row count: the full dataset at N <= 1000, 128 below 5000,
+256 below 50,000, and 512 above. This problem has 20,000 rows, so evolution sees 256
+rows per comparison and the hall of fame is still scored on all 20,000.
+
+Setting `batching=False` runs every population comparison against all 20,000 rows.
+That is the same search on the same data, and it is far slower.
+"""
+
+import numpy as np
+
+from pysr import PySRRegressor
+
+rng = np.random.default_rng(0)
+X = rng.uniform(-3, 3, (20000, 5))
+y = 2.5382 * np.cos(X[:, 3]) + X[:, 0] ** 2 - 0.5
+
+VARIABLE_NAMES = ["x0", "x1", "x2", "x3", "x4"]
+
+MODEL_KWARGS = dict(
+ binary_operators=["+", "-", "*", "/"],
+ unary_operators=["cos", "exp"],
+ batching="auto",
+ niterations=40,
+ verbosity=0,
+)
+
+
+def check(model):
+ """The target is recovered when a front member reaches numerically zero loss."""
+ return bool(model.equations_["loss"].min() < 1e-10)
+
+
+def main():
+ model = PySRRegressor(**MODEL_KWARGS, random_state=0)
+ model.fit(X, y, variable_names=VARIABLE_NAMES)
+ print(model.equations_[["complexity", "loss", "equation"]].to_string(index=False))
+ print(f"\nbest loss: {model.equations_['loss'].min():.3e}")
+ print(f"recovered: {check(model)}")
+
+
+if __name__ == "__main__":
+ main()
diff --git a/examples/backsolve_mutation.py b/examples/backsolve_mutation.py
new file mode 100644
index 000000000..2ea847568
--- /dev/null
+++ b/examples/backsolve_mutation.py
@@ -0,0 +1,123 @@
+"""Recover 2.5*cos(3x) + 0.5x^2 - 1 with the backsolve mutation switched on.
+
+Backsolve inverts the operators sitting above a chosen subtree, which turns
+"improve this whole expression" into "solve for the piece that belongs in this
+hole", then fits that piece as a sparse weighted sum of parts already present in
+the population. It is off by default (`weight_backsolve` maps to weight 0.0);
+this target is out of reach for the ordinary mutations alone.
+"""
+
+import numpy as np
+
+from pysr import PySRRegressor
+
+x = np.linspace(-3.0, 3.0, 200)
+X = x.reshape(-1, 1)
+y = 2.5 * np.cos(3.0 * x) + 0.5 * x * x - 1.0
+VARIABLE_NAMES = ["x"]
+
+MODEL_KWARGS = dict(
+ binary_operators=["+", "-", "*"],
+ unary_operators=["cos", "exp"],
+ weight_backsolve=1.0,
+ maxsize=30,
+ precision=64,
+ niterations=10,
+ verbosity=0,
+)
+
+# One backsolve event, applied outside the search loop so its two losses can be
+# read directly. `mutate!` returns a member only when its monotone gate has
+# already confirmed the child beats the parent on the data.
+_BACKSOLVE_EVENT = """
+function backsolve_event(dataset, population, options, mutation)
+ events = NamedTuple[]
+ for member in population.members
+ tree = copy(member.tree)
+ result = SymbolicRegression.mutate!(
+ tree, member, mutation, options;
+ trace=nothing, dataset, curmaxsize=options.maxsize,
+ nfeatures=size(dataset.X, 1), population_for_backsolve=population,
+ )
+ result.member === nothing && continue
+ push!(events, (
+ before=Float64(member.loss),
+ after=Float64(result.member.loss),
+ parent=SymbolicRegression.string_tree(member.tree, options),
+ child=SymbolicRegression.string_tree(result.member.tree, options),
+ ))
+ end
+ return events
+end
+"""
+
+
+def check(model):
+ """The target recovered to double-precision round-off."""
+ prediction = np.asarray(model.predict(X), dtype=float)
+ return bool(np.max(np.abs(prediction - y)) < 1e-10)
+
+
+def backsolve_events(model, seed=0):
+ """Apply one backsolve mutation to every member of the search's populations.
+
+ Returns the events the mutation's own acceptance gate let through, each with
+ the parent's loss before and the child's loss after.
+ """
+ from pysr.julia_import import SymbolicRegression as SR
+ from pysr.julia_import import jl
+
+ jl.seval("using Random: seed!")
+ jl.seval(f"seed!({seed})")
+ run_event = jl.seval(_BACKSOLVE_EVENT)
+
+ dtype = np.float64 if MODEL_KWARGS["precision"] == 64 else np.float32
+ dataset = SR.Dataset(
+ jl.Matrix(np.asfortranarray(X.T.astype(dtype))), jl.Vector(y.astype(dtype))
+ )
+ mutation = SR.BacksolveMutation()
+ options = model.julia_options_
+ populations = model.julia_state_[0][0]
+
+ events, swept = [], 0
+ for i in range(int(jl.length(populations))):
+ population = populations[i]
+ swept += int(population.n)
+ found = run_event(dataset, population, options, mutation)
+ for k in range(int(jl.length(found))):
+ e = found[k]
+ events.append(
+ dict(
+ before=float(e.before),
+ after=float(e.after),
+ parent=str(e.parent),
+ child=str(e.child),
+ )
+ )
+ return events, swept
+
+
+def main():
+ model = PySRRegressor(**MODEL_KWARGS, random_state=0)
+ model.fit(X, y, variable_names=VARIABLE_NAMES)
+ print(model.equations_[["complexity", "loss", "equation"]].to_string(index=False))
+
+ prediction = np.asarray(model.predict(X), dtype=float)
+ print(f"\nmax abs error: {np.max(np.abs(prediction - y)):.3g}")
+ print(f"recovered: {check(model)}")
+
+ events, swept = backsolve_events(model)
+ print(f"\nbacksolve events accepted: {len(events)} of {swept} members swept")
+ if events:
+ ratios = sorted(e["after"] / e["before"] for e in events if e["before"] > 0)
+ print(f"median child/parent loss ratio: {ratios[len(ratios) // 2]:.3g}")
+ best = min(
+ events, key=lambda e: e["after"] / e["before"] if e["before"] > 0 else 1.0
+ )
+ print(f"largest drop: loss {best['before']:.6g} -> {best['after']:.6g}")
+ print(f" parent: {best['parent']}")
+ print(f" child: {best['child']}")
+
+
+if __name__ == "__main__":
+ main()
diff --git a/examples/cartpole_objective.py b/examples/cartpole_objective.py
new file mode 100644
index 000000000..a1842ebf8
--- /dev/null
+++ b/examples/cartpole_objective.py
@@ -0,0 +1,223 @@
+"""Search for a cart-pole swing-up policy through a custom expression-level objective.
+
+There is no target column. The loss takes a whole candidate expression, drives a cart-pole
+plant with it for ten simulated seconds from each of sixteen starts, and returns minus the
+mean per-step reward. The reward pays for an upright pole and charges for cart offset,
+motor effort and jerk. A rollout that leaves the numerical-divergence box scores a
+sentinel.
+
+Most of the starts hang the pole straight down, so a policy has to pump energy in, catch
+the pole at the top and hold it there, all from one closed-form expression of the five
+observations. The expression returns a force in units of the ten-newton actuator cap.
+
+Reward-shaped losses go negative, which the default logarithmic complexity scaling cannot
+represent, so the search uses linear loss scaling.
+
+Success is the champion earning positive mean reward on all sixty-four held-out starts,
+which are drawn wider than the training starts in every coordinate.
+"""
+
+import numpy as np
+import sympy
+
+from pysr import PySRRegressor
+
+M_C, M_P, LENGTH, GRAVITY = 1.0, 0.1, 0.5, 9.8
+FORCE_CAP, DT = 10.0, 0.02
+RAIL, V_SCALE, OMEGA_SCALE = 2.4, 4.0, 8.0
+HORIZON = 500 # 10 s of control at dt = 0.02
+
+
+def _starts(seed, rows, groups):
+ """Starts drawn one coordinate at a time, in the order the film registered them."""
+ rng = np.random.Generator(np.random.PCG64DXSM(seed))
+ out = [(x, 0.0, np.pi, 0.0) for x in rows]
+ for count, x_limit, centre, theta_limit, omega_limit in groups:
+ for _ in range(count):
+ out.append(
+ (
+ rng.uniform(-x_limit, x_limit),
+ rng.uniform(-x_limit, x_limit),
+ centre + rng.uniform(-theta_limit, theta_limit),
+ rng.uniform(-omega_limit, omega_limit),
+ )
+ )
+ return np.asarray(out, dtype=float)
+
+
+TRAIN_STARTS = _starts(
+ 20260819,
+ (0.0, 0.25),
+ (
+ (6, 0.35, np.pi, 0.35, 0.5),
+ (2, 0.35, np.pi / 2, 0.20, 0.5),
+ (2, 0.35, -np.pi / 2, 0.20, 0.5),
+ (4, 0.20, 0.0, 0.18, 0.5),
+ ),
+)
+HELD_OUT_STARTS = _starts(
+ 20260820,
+ (-0.40, -0.15, 0.15, 0.40),
+ (
+ (28, 0.55, np.pi, 0.70, 1.0),
+ (8, 0.55, np.pi / 2, 0.35, 1.0),
+ (8, 0.55, -np.pi / 2, 0.35, 1.0),
+ (16, 0.30, 0.0, 0.25, 1.0),
+ ),
+)
+
+
+def step(state, force):
+ """One semi-implicit Euler action interval for a batch of cart-pole states."""
+ x, v, theta, omega = state.T
+ total = M_C + M_P
+ sin, cos = np.sin(theta), np.cos(theta)
+ q = (force + M_P * LENGTH * omega**2 * sin) / total
+ theta_dd = (GRAVITY * sin - cos * q) / (LENGTH * (4 / 3 - M_P * cos**2 / total))
+ x_dd = q - M_P * LENGTH * theta_dd * cos / total
+ v, omega = v + DT * x_dd, omega + DT * theta_dd
+ return np.stack([x + DT * v, v, theta + DT * omega, omega], axis=-1)
+
+
+def observe(state):
+ """The five policy inputs: scaled cart offset and speed, the pole as a unit vector,
+ and scaled pole rate. They encode the four physical coordinates without loss."""
+ x, v, theta, omega = state.T
+ return np.stack(
+ [x / RAIL, v / V_SCALE, np.sin(theta), np.cos(theta), omega / OMEGA_SCALE],
+ axis=-1,
+ )
+
+
+X = observe(TRAIN_STARTS)
+y = np.zeros(len(X))
+VARIABLE_NAMES = ["x_n", "v_n", "s", "c", "omega_n"]
+
+CARTPOLE_REWARD = """
+function cartpole_reward(ex, dataset::Dataset{T,L}, options)::L where {T,L}
+ n = size(dataset.X, 2)
+ x = %(rail)r .* Float64.(dataset.X[1, :])
+ v = %(v_scale)r .* Float64.(dataset.X[2, :])
+ theta = atan.(Float64.(dataset.X[3, :]), Float64.(dataset.X[4, :]))
+ omega = %(omega_scale)r .* Float64.(dataset.X[5, :])
+
+ obs = Matrix{Float64}(undef, 5, n)
+ previous = zeros(Float64, n)
+ total = %(m_c)r + %(m_p)r
+ earned = 0.0
+ for _ in 1:%(horizon)d
+ sin_t, cos_t = sin.(theta), cos.(theta)
+ obs[1, :] .= x ./ %(rail)r
+ obs[2, :] .= v ./ %(v_scale)r
+ obs[3, :] .= sin_t
+ obs[4, :] .= cos_t
+ obs[5, :] .= omega ./ %(omega_scale)r
+ raw, ok = eval_tree_array(ex, obs, options)
+ # A batched evaluation is only usable if it returned one force per start.
+ (ok && length(raw) == n && all(isfinite, raw)) || return L(Inf)
+
+ force = %(cap)r .* clamp.(raw, -1.0, 1.0)
+ earned += sum(
+ 2.0 .* cos_t .- 0.05 .* x .^ 2 .-
+ 0.01 .* (force ./ %(cap)r) .^ 2 .-
+ 0.005 .* ((force .- previous) ./ %(cap)r) .^ 2
+ )
+ previous = force
+
+ q = (force .+ %(m_p)r * %(length)r .* omega .^ 2 .* sin_t) ./ total
+ theta_dd = (%(gravity)r .* sin_t .- cos_t .* q) ./
+ (%(length)r .* (4 / 3 .- %(m_p)r .* cos_t .^ 2 ./ total))
+ x_dd = q .- %(m_p)r * %(length)r .* theta_dd .* cos_t ./ total
+ v .+= %(dt)r .* x_dd
+ omega .+= %(dt)r .* theta_dd
+ x .+= %(dt)r .* v
+ theta .+= %(dt)r .* omega
+
+ all(isfinite, x) && all(isfinite, v) && all(isfinite, omega) &&
+ maximum(abs, x) <= 20.0 && maximum(abs, v) <= 100.0 &&
+ maximum(abs, omega) <= 100.0 || return L(1.0e12)
+ end
+ return L(-earned / (%(horizon)d * n))
+end
+""" % dict(
+ m_c=M_C,
+ m_p=M_P,
+ length=LENGTH,
+ gravity=GRAVITY,
+ cap=FORCE_CAP,
+ dt=DT,
+ rail=RAIL,
+ v_scale=V_SCALE,
+ omega_scale=OMEGA_SCALE,
+ horizon=HORIZON,
+)
+
+MODEL_KWARGS = dict(
+ operators={
+ 1: ["square", "abs", "tanh"],
+ 2: ["+", "-", "*", "/", "max", "min"],
+ 3: ["ifelse(t, a, b) = t > 0 ? a : b"],
+ },
+ extra_sympy_mappings={
+ "ifelse": lambda t, a, b: sympy.Piecewise((a, t > 0), (b, True))
+ },
+ loss_function_expression=CARTPOLE_REWARD,
+ loss_scale="linear", # rewards make the loss negative, which log scaling forbids
+ maxsize=28,
+ maxdepth=10,
+ parsimony=0.001,
+ niterations=120,
+ populations=32,
+ population_size=64,
+ verbosity=0,
+)
+
+
+def mean_rewards(model, starts, index=None):
+ """Per-start mean reward over the full horizon, with the training divergence rule."""
+ state = np.array(starts, dtype=float)
+ previous = np.zeros(len(state))
+ earned = np.zeros(len(state))
+ alive = np.ones(len(state), dtype=bool)
+ for _ in range(HORIZON):
+ raw = np.asarray(model.predict(observe(state), index=index), dtype=float)
+ force = FORCE_CAP * np.clip(np.nan_to_num(raw, nan=0.0), -1.0, 1.0)
+ earned += np.where(
+ alive,
+ 2.0 * np.cos(state[:, 2])
+ - 0.05 * state[:, 0] ** 2
+ - 0.01 * (force / FORCE_CAP) ** 2
+ - 0.005 * ((force - previous) / FORCE_CAP) ** 2,
+ 0.0,
+ )
+ previous = force
+ state = step(state, force)
+ alive &= np.isfinite(state).all(axis=-1)
+ alive &= np.abs(state[:, 0]) <= 20.0
+ alive &= np.abs(state[:, 1]) <= 100.0
+ alive &= np.abs(state[:, 3]) <= 100.0
+ state = np.nan_to_num(state, nan=0.0, posinf=0.0, neginf=0.0)
+ return np.where(alive, earned / HORIZON, -1.0e12)
+
+
+def check(model):
+ """The champion earns positive mean reward on every one of the 64 held-out starts."""
+ index = int(model.equations_["loss"].idxmin())
+ return bool((mean_rewards(model, HELD_OUT_STARTS, index=index) > 0.0).all())
+
+
+def main():
+ model = PySRRegressor(**MODEL_KWARGS, random_state=0)
+ model.fit(X, y, variable_names=VARIABLE_NAMES)
+ print(model.equations_[["complexity", "loss", "equation"]].to_string(index=False))
+
+ index = int(model.equations_["loss"].idxmin())
+ row = model.equations_.loc[index]
+ print(f"\nchampion (complexity {row.complexity}, loss {row.loss:+.6f}):")
+ print(f" {row.equation}")
+ returns = mean_rewards(model, HELD_OUT_STARTS, index=index)
+ print(f"held-out positive: {int((returns > 0).sum())}/64 starts")
+
+
+if __name__ == "__main__":
+ main()
diff --git a/examples/conftest.py b/examples/conftest.py
new file mode 100644
index 000000000..06ab2e95d
--- /dev/null
+++ b/examples/conftest.py
@@ -0,0 +1,30 @@
+"""Marker registration for the documentation-example suite.
+
+The slow examples are real searches measured in minutes, so they are opt-in. A bare
+`pytest` run exercises every example that finishes in about a minute.
+"""
+
+
+def pytest_addoption(parser):
+ parser.addoption(
+ "--runslow",
+ action="store_true",
+ default=False,
+ help="also run the examples whose searches take minutes",
+ )
+
+
+def pytest_configure(config):
+ config.addinivalue_line("markers", "slow: a search measured in minutes")
+
+
+def pytest_collection_modifyitems(config, items):
+ if config.getoption("--runslow"):
+ return
+
+ import pytest
+
+ skip = pytest.mark.skip(reason="needs --runslow")
+ for item in items:
+ if "slow" in item.keywords:
+ item.add_marker(skip)
diff --git a/examples/game_of_life.py b/examples/game_of_life.py
new file mode 100644
index 000000000..48cc74fd5
--- /dev/null
+++ b/examples/game_of_life.py
@@ -0,0 +1,58 @@
+"""Rediscover Conway's Game of Life (B3/S23) from its transition table.
+
+The search sees every combination of a cell state and a neighbour count, and nothing
+else. Its vocabulary is and / or / not plus an integer equality test; the neighbour
+thresholds are integer constants it has to find.
+"""
+
+import numpy as np
+
+from pysr import PySRRegressor, TypeSpec
+
+TABLE = [(a, n, int(n == 3 or (a == 1 and n == 2))) for a in (0, 1) for n in range(9)]
+REPLICAS = 8
+X = np.array([[a, n] for a, n, _ in TABLE] * REPLICAS, dtype=object)
+y = np.array([t for _, _, t in TABLE] * REPLICAS, dtype=object)
+VARIABLE_NAMES = ["alive", "n"]
+
+SPEC = TypeSpec(
+ "Cell",
+ fields={"v": "Int"},
+ sample="rng -> Cell(rand(rng, 0:8))",
+ mutate="(rng, value, temperature) -> Cell(mod(value.v + rand(rng, (-1, 1)), 9))",
+ string="value -> string(value.v)",
+)
+
+MODEL_KWARGS = dict(
+ type_spec=SPEC,
+ operators={
+ 1: ["not(a::Cell) = Cell(a.v == 0 ? 1 : 0)"],
+ 2: [
+ "and(a::Cell, b::Cell) = Cell((a.v != 0 && b.v != 0) ? 1 : 0)",
+ "or(a::Cell, b::Cell) = Cell((a.v != 0 || b.v != 0) ? 1 : 0)",
+ "eq(a::Cell, b::Cell) = Cell(a.v == b.v ? 1 : 0)",
+ ],
+ },
+ elementwise_loss="cell_loss(prediction::Cell, target::Cell)::Float64 = prediction.v == target.v ? 0.0 : 1.0",
+ niterations=160,
+ verbosity=0,
+)
+
+
+def check(model):
+ """Every row of the transition table reproduced exactly."""
+ prediction = np.array([int(v) for v in model.predict(X)])
+ return bool((prediction == np.array([int(v) for v in y])).all())
+
+
+def main():
+ model = PySRRegressor(**MODEL_KWARGS, random_state=0)
+ model.fit(X, y, variable_names=VARIABLE_NAMES)
+ print(model.equations_[["complexity", "loss", "equation"]].to_string(index=False))
+ prediction = np.array([int(v) for v in model.predict(X)])
+ target = np.array([int(v) for v in y])
+ print(f"\nexact rows: {int((prediction == target).sum())}/{len(target)}")
+
+
+if __name__ == "__main__":
+ main()
diff --git a/examples/hopper_controller.py b/examples/hopper_controller.py
new file mode 100644
index 000000000..c4da64105
--- /dev/null
+++ b/examples/hopper_controller.py
@@ -0,0 +1,119 @@
+"""Replay the hopping controller PySR evolved against MuJoCo rollouts.
+
+Every candidate expression was scored by simulating the Gymnasium Hopper-v5 and
+accumulating its reward, so the search optimised locomotion directly. There was
+no expert policy and no imitation data. The search is a cluster job; this script
+replays what it found, against the twenty held-out start states that were frozen
+before the search and never scored during it.
+
+Requires `gymnasium[mujoco]`.
+"""
+
+import gymnasium as gym
+import numpy as np
+
+HELDOUT_SEEDS = range(20261001, 20261021)
+RESET_NOISE = 0.01
+
+# Eight seconds at the Hopper's 0.008 s control step. Passing means staying
+# upright for all of it, travelling five metres, and completing four
+# contact/liftoff cycles.
+STEPS = 1000
+DISTANCE = 5.0
+CYCLES = 4
+
+# The three vector constants the search chose, one component per motor.
+C0 = np.array([-0.24771142177432096, -4.22684590927177286, 8.8397066748254275])
+C1 = np.array([0.06452662644513463, 0.6180986305329657, 0.3968791265257595])
+C2 = np.array([0.49735344234321904, -1.1920811500492086, 1.1808251376288732])
+
+
+def action(qpos, qvel):
+ """The evolved controller: hopper state in, three motor commands out."""
+ angles = qpos[3:6]
+ rates = np.clip(qvel[3:6], -10.0, 10.0) / 10.0
+ pitch = qpos[2]
+ forward, vertical, spin = qvel[0] / 5.0, qvel[1] / 5.0, qvel[2] / 5.0
+ rolled_angles = angles[[1, 2, 0]]
+ rolled_rates = rates[[1, 2, 0]]
+ d = angles - 2 * rolled_angles + rolled_rates - 2 * (forward + spin) + C0
+ left = -angles - (pitch + spin * (rolled_rates - spin)) * d - C1
+ right = -rolled_angles - vertical + spin + C2
+ return np.tanh(rolled_angles - spin + left * right)
+
+
+def contact_geoms(env):
+ model = env.unwrapped.model
+ named = [(g, (model.geom(g).name or "").lower()) for g in range(model.ngeom)]
+ return (
+ {g for g, name in named if "foot" in name},
+ {g for g, name in named if "floor" in name},
+ )
+
+
+def touching(data, feet, floor):
+ pairs = ((int(c.geom1), int(c.geom2)) for c in data.contact[: data.ncon])
+ return any(a in feet and b in floor or b in feet and a in floor for a, b in pairs)
+
+
+def cycles(contact):
+ """Completed liftoff-to-landing cycles, which is what makes a gait a gait."""
+ count = 0
+ landed = airborne = False
+ for down in contact:
+ if down:
+ count += airborne
+ landed, airborne = True, False
+ elif landed:
+ airborne = True
+ return count
+
+
+def trial(env, feet, floor, seed):
+ env.reset(seed=seed)
+ data = env.unwrapped.data
+ start = float(data.qpos[0])
+ contact = []
+ upright = True
+ for _ in range(STEPS):
+ _, _, terminated, truncated, _ = env.step(action(data.qpos, data.qvel))
+ contact.append(touching(data, feet, floor))
+ upright = not terminated
+ if terminated or truncated:
+ break
+ hops = cycles(contact)
+ distance = float(data.qpos[0]) - start
+ return {
+ "seed": seed,
+ "steps": len(contact),
+ "distance": distance,
+ "cycles": hops,
+ "success": (
+ len(contact) == STEPS
+ and upright
+ and distance >= DISTANCE
+ and hops >= CYCLES
+ ),
+ }
+
+
+def replay():
+ env = gym.make("Hopper-v5", reset_noise_scale=RESET_NOISE)
+ feet, floor = contact_geoms(env)
+ try:
+ return [trial(env, feet, floor, seed) for seed in HELDOUT_SEEDS]
+ finally:
+ env.close()
+
+
+def main():
+ results = replay()
+ for r in results:
+ outcome = "passed" if r["success"] else "fell"
+ line = f"seed {r['seed']} {outcome:<6} {r['distance']:6.2f} m"
+ print(f"{line} {r['cycles']:2d} hops {r['steps']:4d} steps")
+ print(f"\n{sum(r['success'] for r in results)} of {len(results)} held-out starts")
+
+
+if __name__ == "__main__":
+ main()
diff --git a/examples/magnetic_field.py b/examples/magnetic_field.py
new file mode 100644
index 000000000..21ca31c16
--- /dev/null
+++ b/examples/magnetic_field.py
@@ -0,0 +1,146 @@
+"""Recover a magnetic field from force measurements alone.
+
+A charged particle obeys F = -eta(T) v + v x B(t), and 1000 experiments record only
+the total force and the inputs (t, v, T). A TemplateExpressionSpec supplies that
+structure and a Force type spec carries the three-component data, so the search has
+only to fill in B_x(t), B_y(t), B_z(t) and the drag scale eta(T).
+"""
+
+import numpy as np
+
+from pysr import PySRRegressor, TemplateExpressionSpec, TypeSpec
+
+OMEGA = 2 * np.pi
+
+
+def experiments(n, seed):
+ rng = np.random.default_rng(seed)
+ T = 298.15 + 0.5 * rng.random(n)
+ t = 10.0 * rng.random(n)
+ v = 2.0 * rng.random((n, 3)) - 1.0
+ B = np.stack([np.sin(OMEGA * t), np.cos(OMEGA * t), np.exp(-t / 10.0)], axis=1)
+ F = -1e-5 * np.sqrt(T)[:, None] * v + np.cross(v, B)
+ return np.column_stack([t, v[:, 0], v[:, 1], v[:, 2], T]), F
+
+
+def as_forces(rows):
+ """An (n, 3) float array as a column of n Force values."""
+ out = np.empty(len(rows), dtype=object)
+ out[:] = [tuple(map(float, row)) for row in rows]
+ return out
+
+
+def on_diagonal(columns):
+ """Every value in the problem is a Force, so scalar inputs ride the diagonal."""
+ out = np.empty(columns.shape, dtype=object)
+ for j in range(columns.shape[1]):
+ out[:, j] = as_forces(np.repeat(columns[:, j, None], 3, axis=1))
+ return out
+
+
+INPUTS, FORCES = experiments(1000, seed=0)
+X = on_diagonal(INPUTS)
+y = as_forces(FORCES)
+VARIABLE_NAMES = ["t", "v_x", "v_y", "v_z", "T"]
+
+HELD_OUT_INPUTS, HELD_OUT_FORCES = experiments(500, seed=12345)
+HELD_OUT_X = on_diagonal(HELD_OUT_INPUTS)
+
+FORCE = TypeSpec(
+ "Force",
+ fields={"x": "Float64", "y": "Float64", "z": "Float64"},
+ sample="rng -> Force(randn(rng, 3)...)",
+ definitions="""
+ Base.sin(a::Force)::Force = Force(sin(a.x), sin(a.y), sin(a.z))
+ Base.cos(a::Force)::Force = Force(cos(a.x), cos(a.y), cos(a.z))
+ function Base.sqrt(a::Force)::Force
+ f(v) = v < 0 ? NaN : sqrt(v)
+ return Force(f(a.x), f(a.y), f(a.z))
+ end
+ Base.exp(a::Force)::Force = Force(exp(a.x), exp(a.y), exp(a.z))
+ Base.:+(a::Force, b::Force)::Force = Force(a.x + b.x, a.y + b.y, a.z + b.z)
+ Base.:-(a::Force, b::Force)::Force = Force(a.x - b.x, a.y - b.y, a.z - b.z)
+ Base.:*(a::Force, b::Force)::Force = Force(a.x * b.x, a.y * b.y, a.z * b.z)
+ Base.:/(a::Force, b::Force)::Force = Force(a.x / b.x, a.y / b.y, a.z / b.z)
+ """,
+ scalar_constants="value -> [value.x, value.y, value.z]",
+ with_scalar_constants="(value, c) -> Force(c[1], c[2], c[3])",
+ string="value -> sprint(show, (value.x, value.y, value.z); context = :compact => true)",
+ init_invalid="() -> Force(NaN, NaN, NaN)",
+)
+
+
+OPERATORS = {
+ 1: [
+ "Base.sin",
+ "Base.cos",
+ "Base.sqrt",
+ "Base.exp",
+ ],
+ 2: [
+ "Base.:+",
+ "Base.:-",
+ "Base.:*",
+ "Base.:/",
+ ],
+}
+
+FORCE_LOSS = "force_loss(a::Force, b::Force)::Float64 = (a.x - b.x)^2 + (a.y - b.y)^2 + (a.z - b.z)^2"
+
+PHYSICS = r"""
+begin
+ _B_x = B_x(t)
+ _B_y = B_y(t)
+ _B_z = B_z(t)
+ _F_d_scale = F_d_scale(T)
+ if !(_B_x.valid && _B_y.valid && _B_z.valid && _F_d_scale.valid)
+ return ValidVector(_B_x.x, false)
+ end
+ F = map(_B_x.x, _B_y.x, _B_z.x, _F_d_scale.x,
+ v_x.x, v_y.x, v_z.x) do bx, by, bz, fd, ux, uy, uz
+ b1, b2, b3 = bx.x, by.x, bz.x
+ u1, u2, u3 = ux.x, uy.x, uz.x
+ s = fd.x
+ Force(u2 * b3 - u3 * b2 + s * u1,
+ u3 * b1 - u1 * b3 + s * u2,
+ u1 * b2 - u2 * b1 + s * u3)
+ end
+ ValidVector(F, true)
+end
+"""
+
+STRUCTURE = TemplateExpressionSpec(
+ combine=PHYSICS,
+ expressions=["B_x", "B_y", "B_z", "F_d_scale"],
+ variable_names=VARIABLE_NAMES,
+)
+
+MODEL_KWARGS = dict(
+ type_spec=FORCE,
+ expression_spec=STRUCTURE,
+ operators=OPERATORS,
+ elementwise_loss=FORCE_LOSS,
+ niterations=100,
+ verbosity=0,
+)
+
+
+def held_out_error(model):
+ predicted = np.array([list(f) for f in model.predict(HELD_OUT_X)], dtype=float)
+ return float(np.abs(predicted - HELD_OUT_FORCES).max())
+
+
+def check(model):
+ """Forces on unseen experiments reproduced to machine precision."""
+ return held_out_error(model) < 1e-6
+
+
+def main():
+ model = PySRRegressor(**MODEL_KWARGS, random_state=0)
+ model.fit(X, y, variable_names=VARIABLE_NAMES)
+ print(model.equations_[["complexity", "loss", "equation"]].to_string(index=False))
+ print(f"\nheld-out max abs error: {held_out_error(model):.3g}")
+
+
+if __name__ == "__main__":
+ main()
diff --git a/examples/morton_interleave.py b/examples/morton_interleave.py
new file mode 100644
index 000000000..87f658b44
--- /dev/null
+++ b/examples/morton_interleave.py
@@ -0,0 +1,114 @@
+"""Rediscover a bit-exact Morton (Z-order) interleave of two 4-bit integers.
+
+The search works on a machine word: a Word value backed by UInt32 whose vocabulary is
+the integer instruction set (xor, and, or, not, shift left, shift right), with shift
+distances as ordinary Word constants it has to find. It scores all 256 four-bit input
+pairs and has to recover the map that interleaves their bits.
+
+The score only ever sees those 256 pairs, so what comes back is an expression that
+agrees with the interleave everywhere it was asked. That is a weaker claim than the
+general Z-order map, and the difference is measurable: the recovered forms are exact on
+every 4-bit pair and correct on only 256 of the 1024 5-bit pairs, which are exactly the
+4-bit ones. The run prints both numbers.
+
+Two things about the shape of the answer are worth knowing before reading it. Written in
+this vocabulary as a tree, with no shared subexpressions, the textbook interleave kernel
+(two shift-mask spread stages per input, then one shift and or) is 41 nodes, so a
+recovered form near 40 is the same size as the hand-written one rather than bloat, and
+maxsize has to leave room for it. And the masks print as arbitrary hex because only
+their low bits can affect an 8-bit result: in 0x00080022 the search pinned bits 1 and 5
+and bit 19 is free, since no scored pair can tell what it holds.
+"""
+
+import numpy as np
+
+from pysr import PySRRegressor, TypeSpec
+
+WIDTH = 4
+WIDER = 5
+
+
+def morton(x, y, width=WIDTH):
+ z = 0
+ for i in range(width):
+ z |= ((x >> i) & 1) << (2 * i)
+ z |= ((y >> i) & 1) << (2 * i + 1)
+ return z
+
+
+PAIRS = [(a, b) for b in range(1 << WIDTH) for a in range(1 << WIDTH)]
+X = np.array(PAIRS, dtype=object)
+y = np.array([morton(a, b) for a, b in PAIRS], dtype=object)
+VARIABLE_NAMES = ["x", "y"]
+
+SPEC = TypeSpec(
+ "Word",
+ fields={"bits": "UInt32"},
+ sample="rng -> Word(rand(rng, Bool) ? rand(rng, UInt32(0):UInt32(31)) : rand(rng, UInt32))",
+ mutate="(rng, value, temperature) -> Word(rand(rng, Bool) ? xor(value.bits, UInt32(1) << rand(rng, 0:31)) : value.bits + rand(rng, (UInt32(1), typemax(UInt32))))",
+ string='value -> value.bits < UInt32(32) ? string(value.bits) : "0x" * string(value.bits, base = 16, pad = 8)',
+)
+
+MODEL_KWARGS = dict(
+ type_spec=SPEC,
+ operators={
+ 1: ["bnot(a::Word) = Word(~a.bits)"],
+ 2: [
+ "bxor(a::Word, b::Word) = Word(xor(a.bits, b.bits))",
+ "band(a::Word, b::Word) = Word(a.bits & b.bits)",
+ "bor(a::Word, b::Word) = Word(a.bits | b.bits)",
+ "shl(a::Word, b::Word) = Word(a.bits << b.bits)",
+ "shr(a::Word, b::Word) = Word(a.bits >> b.bits)",
+ ],
+ },
+ elementwise_loss="bit_loss(prediction::Word, target::Word)::Float64 = count_ones(xor(prediction.bits, target.bits)) / 32",
+ maxsize=45,
+ niterations=800,
+ verbosity=0,
+)
+
+
+def exact_front_rows(model):
+ """Indices of front members that reproduce every one of the 256 interleaves."""
+ target = np.array([int(v) for v in y])
+ rows = []
+ for i in range(len(model.equations_)):
+ prediction = np.array([int(v) for v in model.predict(X, index=i)])
+ if (prediction == target).all():
+ rows.append(i)
+ return rows
+
+
+def check(model):
+ """Success is an exact interleaver anywhere on the Pareto front."""
+ return bool(exact_front_rows(model))
+
+
+def wider_agreement(model, index):
+ """How many 5-bit pairs the recovered expression interleaves correctly."""
+ pairs = [(a, b) for b in range(1 << WIDER) for a in range(1 << WIDER)]
+ inputs = np.array(pairs, dtype=object)
+ expected = np.array([morton(a, b, WIDER) for a, b in pairs])
+ got = np.array([int(v) for v in model.predict(inputs, index=index)])
+ return int((got == expected).sum()), len(pairs)
+
+
+def main():
+ model = PySRRegressor(**MODEL_KWARGS, random_state=0)
+ model.fit(X, y, variable_names=VARIABLE_NAMES)
+ print(model.equations_[["complexity", "loss", "equation"]].to_string(index=False))
+ rows = exact_front_rows(model)
+ if rows:
+ best = model.equations_.iloc[rows[0]]
+ print(f"\nexact interleaver at complexity {int(best.complexity)}:")
+ print(best.equation)
+ agree, total = wider_agreement(model, rows[0])
+ print(
+ f"\non {WIDER}-bit pairs it interleaves {agree}/{total} correctly, so it is a {WIDTH}-bit interleaver and not the Z-order map"
+ )
+ else:
+ print("\nno exact interleaver on the front")
+
+
+if __name__ == "__main__":
+ main()
diff --git a/examples/mutations_and_plugins.py b/examples/mutations_and_plugins.py
new file mode 100644
index 000000000..5fb5008e6
--- /dev/null
+++ b/examples/mutations_and_plugins.py
@@ -0,0 +1,64 @@
+"""Weight a mutation and add a plugin from PySRRegressor.
+
+Mutations and plugins are ordinary configuration: `mutations` maps a mutation
+configuration to a weight, `plugins` lists plugin configurations, and entries of
+either override the shipped defaults by type and leave the rest alone.
+`BacksolveMutation` ships at weight 0.0, so giving it a weight is what turns it
+on for this search for 2.5*cos(3x) + 0.5*x^2 - 1.
+"""
+
+import numpy as np
+
+from pysr import AdaptiveMutationWeightsPlugin, BacksolveMutation, PySRRegressor
+from pysr.julia_import import jl
+
+x = np.linspace(-3.0, 3.0, 200)
+X = x.reshape(-1, 1)
+y = 2.5 * np.cos(3.0 * x) + 0.5 * x * x - 1.0
+VARIABLE_NAMES = ["x"]
+
+BASE_KWARGS = dict(
+ binary_operators=["+", "-", "*"],
+ unary_operators=["cos", "exp"],
+ niterations=20,
+ verbosity=0,
+)
+
+MODEL_KWARGS = dict(
+ **BASE_KWARGS,
+ mutations={BacksolveMutation(): 0.1},
+ plugins=[AdaptiveMutationWeightsPlugin()],
+)
+
+
+def check(model):
+ """The noiseless target reproduced to within 1e-5 everywhere."""
+ return bool(np.max(np.abs(model.predict(X) - y)) < 1e-5)
+
+
+def _backsolve_weight(model):
+ for mutation, weight in model.julia_options_.mutations:
+ if "BacksolveMutation" in jl.string(mutation):
+ return float(weight)
+ raise AssertionError("BacksolveMutation missing from the resolved mutation table")
+
+
+def _run(kwargs, label):
+ model = PySRRegressor(**kwargs, random_state=0)
+ model.fit(X, y, variable_names=VARIABLE_NAMES)
+ print(f"=== {label}")
+ print(model.equations_[["complexity", "loss", "equation"]].to_string(index=False))
+ print(f"max error: {np.max(np.abs(model.predict(X) - y)):.3g}")
+ print(f"backsolve weight: {_backsolve_weight(model)}")
+ print(f"plugins: {jl.string(model.julia_options_.plugins)}\n")
+ return model
+
+
+def main():
+ _run(BASE_KWARGS, "default mutations and plugins")
+ model = _run(MODEL_KWARGS, "backsolve weighted, adaptive weights on")
+ print(f"exact recovery: {check(model)}")
+
+
+if __name__ == "__main__":
+ main()
diff --git a/examples/prng_period.py b/examples/prng_period.py
new file mode 100644
index 000000000..a266bc7fd
--- /dev/null
+++ b/examples/prng_period.py
@@ -0,0 +1,466 @@
+"""Invent a 32-bit pseudorandom generator with no target array.
+
+`y` is a column of zeros and the objective never reads it. The search sees 32 seed
+states, and the loss runs each candidate as an iterated 32-bit map and scores the
+generator that map is: how long its orbit is, how balanced its bit columns are, how
+uncorrelated successive states are, how far one flipped input bit spreads, and whether
+any bit row of the output is a shifted copy of a bit row of the input. Every term is 0
+exactly when its desideratum is met.
+
+The period term is the interesting one, because the whole state space is 4294967295
+nonzero words and nothing may walk it. The objective probes the candidate on 0, on the
+32 basis states and on 32 fixed states; if it is GF(2)-linear on that evidence, it
+builds the minimal polynomial of the first seed's Krylov sequence, and if that
+polynomial is irreducible the orbit's length is exactly the multiplicative order of x
+modulo it, obtained by dividing the prime factors out of 2^deg - 1. A map whose period
+cannot be certified is credited only with the orbit actually walked, so a candidate is
+paid for the period it can prove.
+
+`check` re-derives the claim in Python from the recovered expression, by a different
+route than the loss uses. It reads the map's GF(2) matrix out of `predict`, confirms
+the map really is linear on random states, and computes the multiplicative order of
+that matrix by repeated squaring. An order of 4294967295 forces the period: the order
+of a matrix over GF(2) divides the unit-group order of the algebra it generates, which
+is a power of two times a product of 2^d - 1 over the degrees of the minimal
+polynomial's irreducible factors, and 65537 divides 2^d - 1 only when 32 divides d. So
+the minimal polynomial is irreducible of degree 32, the state space is one copy of
+GF(2^32) under the map, and every nonzero word lies on one orbit of length 4294967295.
+"""
+
+import numpy as np
+
+from pysr import PySRRegressor, TypeSpec
+
+WIDTH = 32
+MASK = (1 << WIDTH) - 1
+FULL_PERIOD = MASK
+ORDER_PRIMES = (3, 5, 17, 257, 65537)
+
+SEEDS = [int(v) for v in np.random.default_rng(0).integers(1, 1 << WIDTH, size=32)]
+X = np.empty((len(SEEDS), 1), dtype=object)
+for i, seed in enumerate(SEEDS):
+ X[i, 0] = seed
+y = np.zeros(len(SEEDS), dtype=object)
+VARIABLE_NAMES = ["x"]
+
+SPEC = TypeSpec(
+ "Word",
+ fields={"bits": "UInt32"},
+ sample="rng -> Word(rand(rng, Bool) ? rand(rng, UInt32(0):UInt32(31)) : rand(rng, UInt32))",
+ mutate="(rng, value, temperature) -> Word(rand(rng, Bool) ? xor(value.bits, UInt32(1) << rand(rng, 0:31)) : value.bits + rand(rng, (UInt32(1), typemax(UInt32))))",
+ string='value -> value.bits < UInt32(32) ? string(value.bits) : "0x" * string(value.bits, base = 16, pad = 8)',
+ loss_type="Float64",
+)
+
+# The certificate: linear algebra over GF(2), so a period of four billion states costs
+# a few dozen 32-bit operations instead of four billion evaluations.
+PERIOD_CERTIFICATE = r"""
+const PRNG_PROBES = UInt32[UInt32(0x9e3779b9) * UInt32(k) for k in 1:32]
+
+function _prng_matvec(cols::Vector{UInt32}, v::UInt32)
+ acc = UInt32(0)
+ vv = v
+ while vv != 0
+ acc ⊻= cols[trailing_zeros(vv) + 1]
+ vv &= vv - UInt32(1)
+ end
+ return acc
+end
+
+function _prng_mulmod(a::UInt64, b::UInt64, f::UInt64, deg::Int)
+ out = UInt64(0)
+ aa, bb = a, b
+ while bb != 0
+ if bb & 1 == 1
+ out ⊻= aa
+ end
+ bb >>= 1
+ aa <<= 1
+ if (aa >> deg) & 1 == 1
+ aa ⊻= f
+ end
+ end
+ return out
+end
+
+function _prng_powmod(base::UInt64, e::UInt64, f::UInt64, deg::Int)
+ r = UInt64(1)
+ b, ee = base, e
+ while ee != 0
+ if ee & 1 == 1
+ r = _prng_mulmod(r, b, f, deg)
+ end
+ b = _prng_mulmod(b, b, f, deg)
+ ee >>= 1
+ end
+ return r
+end
+
+_prng_deg(p::UInt64) = 63 - leading_zeros(p)
+
+function _prng_gcd(a::UInt64, b::UInt64)
+ x, y = a, b
+ while y != 0
+ dy = _prng_deg(y)
+ while x != 0 && _prng_deg(x) >= dy
+ x ⊻= y << (_prng_deg(x) - dy)
+ end
+ x, y = y, x
+ end
+ return x
+end
+
+function _prng_minpoly(cols::Vector{UInt32}, v0::UInt32)
+ leads = Int[]
+ vecs = UInt32[]
+ combs = UInt64[]
+ cur = v0
+ comb = UInt64(1)
+ for _ in 0:32
+ red, rc = cur, comb
+ for i in eachindex(leads)
+ if (red >> leads[i]) & UInt32(1) == 1
+ red ⊻= vecs[i]
+ rc ⊻= combs[i]
+ end
+ end
+ if red == 0
+ return rc
+ end
+ push!(leads, 31 - leading_zeros(red))
+ push!(vecs, red)
+ push!(combs, rc)
+ cur = _prng_matvec(cols, cur)
+ comb <<= 1
+ end
+ return UInt64(0)
+end
+
+function _prng_irreducible(f::UInt64, deg::Int)
+ deg < 1 && return false
+ _prng_powmod(UInt64(2), UInt64(1) << deg, f, deg) == UInt64(2) || return false
+ d = deg
+ p = 2
+ while p * p <= d
+ if d % p == 0
+ g = _prng_gcd(
+ _prng_powmod(UInt64(2), UInt64(1) << (deg ÷ p), f, deg) ⊻ UInt64(2), f)
+ g == UInt64(1) || return false
+ while d % p == 0
+ d ÷= p
+ end
+ end
+ p += 1
+ end
+ if d > 1
+ g = _prng_gcd(
+ _prng_powmod(UInt64(2), UInt64(1) << (deg ÷ d), f, deg) ⊻ UInt64(2), f)
+ g == UInt64(1) || return false
+ end
+ return true
+end
+
+function _prng_order(f::UInt64, deg::Int)
+ order = (UInt64(1) << deg) - UInt64(1)
+ rest = order
+ p = UInt64(2)
+ while p * p <= rest
+ if rest % p == 0
+ while order % p == 0 && _prng_powmod(UInt64(2), order ÷ p, f, deg) == UInt64(1)
+ order ÷= p
+ end
+ while rest % p == 0
+ rest ÷= p
+ end
+ end
+ p += UInt64(1)
+ end
+ if rest > 1
+ while order % rest == 0 &&
+ _prng_powmod(UInt64(2), order ÷ rest, f, deg) == UInt64(1)
+ order ÷= rest
+ end
+ end
+ return order
+end
+
+function _prng_certified(tree, options, seed::UInt32)
+ m = 33 + length(PRNG_PROBES)
+ grid = Matrix{Word}(undef, 1, m)
+ grid[1, 1] = Word(UInt32(0))
+ for b in 0:31
+ grid[1, 2 + b] = Word(UInt32(1) << b)
+ end
+ for (i, p) in enumerate(PRNG_PROBES)
+ grid[1, 33 + i] = Word(p)
+ end
+ out, ok = eval_tree_array(tree, grid, options)
+ ok || return UInt64(0)
+ out[1].bits == 0 || return UInt64(0)
+
+ cols = [out[2 + b].bits for b in 0:31]
+ for (i, p) in enumerate(PRNG_PROBES)
+ _prng_matvec(cols, p) == out[33 + i].bits || return UInt64(0)
+ end
+
+ f = _prng_minpoly(cols, seed)
+ f == 0 && return UInt64(0)
+ deg = _prng_deg(f)
+ _prng_irreducible(f, deg) || return UInt64(0)
+ return _prng_order(f, deg)
+end
+"""
+
+PRNG_LOSS = PERIOD_CERTIFICATE + r"""
+function prng_loss(tree, dataset::Dataset{T,L}, options)::L where {T,L}
+ steps = 256
+ lags = (1, 2, 3, 5)
+ n = dataset.n
+ walk = Matrix{UInt32}(undef, steps, n)
+ x = copy(dataset.X)
+ for s in 1:steps
+ out, ok = eval_tree_array(tree, x, options)
+ ok || return L(Inf)
+ length(out) == n || error("prng_loss scored " * string(length(out)) *
+ " of " * string(n) * " samples; set batching=false")
+ for j in 1:n
+ walk[s, j] = out[j].bits
+ end
+ x = reshape(out, 1, n)
+ end
+
+ seen = typemax(Int)
+ for j in 1:n
+ seen = min(seen, length(Set(@view walk[:, j])))
+ end
+ certified = _prng_certified(tree, options, dataset.X[1, 1].bits)
+ proven = certified > 0 ? Float64(certified) : Float64(seen)
+ period = 1 - log2(proven) / 32
+
+ ones = zeros(Int, 32)
+ for j in 1:n, s in 1:steps
+ w = walk[s, j]
+ for b in 1:32
+ ones[b] += (w >> (b - 1)) & UInt32(1)
+ end
+ end
+ total = steps * n
+ balance = 0.0
+ for b in 1:32
+ p = ones[b] / total
+ balance += (2p - 1)^2
+ end
+ balance /= 32
+
+ autocorr = 0.0
+ joint = zeros(Int, 32)
+ head = zeros(Int, 32)
+ tail = zeros(Int, 32)
+ for lag in lags
+ fill!(joint, 0)
+ fill!(head, 0)
+ fill!(tail, 0)
+ for j in 1:n, s in 1:(steps - lag)
+ u = walk[s, j]
+ v = walk[s + lag, j]
+ for b in 1:32
+ a = (u >> (b - 1)) & UInt32(1)
+ c = (v >> (b - 1)) & UInt32(1)
+ head[b] += a
+ tail[b] += c
+ joint[b] += a & c
+ end
+ end
+ pairs = (steps - lag) * n
+ for b in 1:32
+ ph = head[b] / pairs
+ pt = tail[b] / pairs
+ cov = joint[b] / pairs - ph * pt
+ var = ph * (1 - ph) * pt * (1 - pt)
+ autocorr += var > 0 ? cov^2 / var : 0.0
+ end
+ end
+ autocorr /= 32 * length(lags)
+
+ m = length(PRNG_PROBES)
+ grid = Matrix{Word}(undef, 1, m)
+ for i in 1:m
+ grid[1, i] = Word(PRNG_PROBES[i])
+ end
+ base, bok = eval_tree_array(tree, grid, options)
+ bok || return L(Inf)
+ moved = 0
+ deps = zeros(Int, 32)
+ for b in 0:31
+ for i in 1:m
+ grid[1, i] = Word(PRNG_PROBES[i] ⊻ (UInt32(1) << b))
+ end
+ flipped, fok = eval_tree_array(tree, grid, options)
+ fok || return L(Inf)
+ touched = UInt32(0)
+ for i in 1:m
+ d = base[i].bits ⊻ flipped[i].bits
+ moved += count_ones(d)
+ touched |= d
+ end
+ for j in 1:32
+ deps[j] += (touched >> (j - 1)) & UInt32(1)
+ end
+ end
+ diffusion = max(0.0, (4.0 - moved / (32 * m)) / 4.0)^2
+ edge = 0.0
+ for j in 1:32
+ edge += max(0.0, (2.0 - deps[j]) / 2.0)^2
+ end
+ edge /= 32
+
+ shear = 0.0
+ npairs = 0
+ sjoint = zeros(Int, 32)
+ shead = zeros(Int, 32)
+ stail = zeros(Int, 32)
+ for lag in (1, 2), d in (-3, -2, -1, 1, 2, 3)
+ blo = d >= 0 ? 0 : -d
+ bhi = d >= 0 ? 31 - d : 31
+ fill!(sjoint, 0)
+ fill!(shead, 0)
+ fill!(stail, 0)
+ for j in 1:n, s in 1:(steps - lag)
+ u = walk[s, j]
+ v = walk[s + lag, j]
+ for b in blo:bhi
+ a = (u >> b) & UInt32(1)
+ c = (v >> (b + d)) & UInt32(1)
+ shead[b + 1] += a
+ stail[b + 1] += c
+ sjoint[b + 1] += a & c
+ end
+ end
+ spairs = (steps - lag) * n
+ for b in blo:bhi
+ ph = shead[b + 1] / spairs
+ pt = stail[b + 1] / spairs
+ cov = sjoint[b + 1] / spairs - ph * pt
+ var = ph * (1 - ph) * pt * (1 - pt)
+ shear += var > 0 ? cov^2 / var : 0.0
+ npairs += 1
+ end
+ end
+ shear /= npairs
+
+ return L(period + balance + autocorr + diffusion + shear + edge)
+end
+"""
+
+MODEL_KWARGS = dict(
+ type_spec=SPEC,
+ operators={
+ 1: ["bnot(a::Word) = Word(~a.bits)"],
+ 2: [
+ "bxor(a::Word, b::Word) = Word(xor(a.bits, b.bits))",
+ "band(a::Word, b::Word) = Word(a.bits & b.bits)",
+ "bor(a::Word, b::Word) = Word(a.bits | b.bits)",
+ "shl(a::Word, b::Word) = Word(a.bits << b.bits)",
+ "shr(a::Word, b::Word) = Word(a.bits >> b.bits)",
+ ],
+ },
+ loss_function=PRNG_LOSS,
+ # The objective walks every seed and indexes the sample axis, so it has to see the
+ # whole sample set; a minibatch would score a different problem.
+ batching=False,
+ maxsize=20,
+ # Eight small populations, evolved deeply: each candidate costs a 256-step walk of
+ # every seed, so the search is worth more generations rather than more members.
+ populations=8,
+ population_size=30,
+ ncycles_per_iteration=30,
+ niterations=800,
+ verbosity=0,
+)
+
+IDENTITY = [1 << b for b in range(WIDTH)]
+LINEARITY_PROBES = [
+ int(v) for v in np.random.default_rng(1).integers(0, 1 << WIDTH, size=256)
+]
+
+
+def matvec(cols, v):
+ """Image of `v` under the GF(2) matrix whose columns are `cols`."""
+ acc = 0
+ while v:
+ acc ^= cols[(v & -v).bit_length() - 1]
+ v &= v - 1
+ return acc
+
+
+def matmul(a, b):
+ return [matvec(a, col) for col in b]
+
+
+def matpow(m, e):
+ result = IDENTITY
+ while e:
+ if e & 1:
+ result = matmul(result, m)
+ m = matmul(m, m)
+ e >>= 1
+ return result
+
+
+def linear_matrix(model, index):
+ """The map's GF(2) matrix, or None if the recovered expression is not linear."""
+ probe = np.empty((1 + WIDTH + len(LINEARITY_PROBES), 1), dtype=object)
+ probe[0, 0] = 0
+ for b in range(WIDTH):
+ probe[1 + b, 0] = 1 << b
+ for i, state in enumerate(LINEARITY_PROBES):
+ probe[1 + WIDTH + i, 0] = state
+ image = [int(v) & MASK for v in model.predict(probe, index=index)]
+ if image[0] != 0:
+ return None
+ cols = image[1 : 1 + WIDTH]
+ for state, got in zip(LINEARITY_PROBES, image[1 + WIDTH :]):
+ if matvec(cols, state) != got:
+ return None
+ return cols
+
+
+def period(model, index):
+ """4294967295 when the map is a full-period generator, 0 when it is not.
+
+ The map is linear, so its order under composition is settled by matrix powers:
+ order 4294967295 forces the minimal polynomial to be irreducible of degree 32, and
+ then every nonzero word sits on a single orbit of that length.
+ """
+ cols = linear_matrix(model, index)
+ if cols is None or matpow(cols, FULL_PERIOD) != IDENTITY:
+ return 0
+ if any(matpow(cols, FULL_PERIOD // q) == IDENTITY for q in ORDER_PRIMES):
+ return 0
+ return FULL_PERIOD
+
+
+def check(model):
+ """Some row of the front runs through all 4294967295 nonzero words."""
+ return any(period(model, i) == FULL_PERIOD for i in model.equations_.index)
+
+
+def main():
+ model = PySRRegressor(**MODEL_KWARGS, random_state=0)
+ model.fit(X, y, variable_names=VARIABLE_NAMES)
+ front = model.equations_[["complexity", "loss", "equation"]].copy()
+ front["full_period"] = [
+ period(model, i) == FULL_PERIOD for i in model.equations_.index
+ ]
+ print(front.to_string(index=False))
+ full = front[front["full_period"]]
+ if len(full):
+ row = full.iloc[0]
+ print(
+ f"\nperiod {FULL_PERIOD} at complexity {int(row.complexity)}: {row.equation}"
+ )
+ print(f"full period reached: {bool(len(full))}")
+
+
+if __name__ == "__main__":
+ main()
diff --git a/examples/pysr_demo.ipynb b/examples/pysr_demo.ipynb
index fd7c264b7..2f2fa51bb 100644
--- a/examples/pysr_demo.ipynb
+++ b/examples/pysr_demo.ipynb
@@ -138,7 +138,7 @@
"id": "N4gANfkaj8ie"
},
"source": [
- "PySR can run for arbitrarily long, and continue to find more and more accurate expressions. You can set the total number of cycles of evolution with `niterations`, although there are also a [few more ways](https://github.com/MilesCranmer/PySR/pull/134) to stop execution.\n",
+ "PySR can run for arbitrarily long, and continue to find more and more accurate expressions. You can set the total number of cycles of evolution with `niterations`, although there are also a [few more ways](https://github.com/astroautomata/PySR/pull/134) to stop execution.\n",
"\n",
"**This first execution will take a bit longer to startup, as the library is JIT-compiled. The next execution will be much faster.**"
]
@@ -321,7 +321,7 @@
"id": "qvgVbOoSFtQY"
},
"source": [
- "A full list of operators is given here: https://ai.damtp.cam.ac.uk/pysr/operators,\n",
+ "A full list of operators is given here: https://pysr.ai/operators,\n",
"but we can also use any binary or unary operator in `julia`, or define our own as arbitrary functions.\n",
"\n",
"Say that we want a command to do quartic powers:\n",
@@ -801,8 +801,8 @@
"cell_type": "markdown",
"metadata": {},
"source": [
- "Note that PySR should automatically load the `juliacall.ipython` extension for you,\n",
- "which means that you can also execute Julia code in the notebook using the `%%julia` magic:"
+ "You can also execute Julia code in the notebook using the `%%julia` magic.\n",
+ "Load the `juliacall` IPython extension to enable it:"
]
},
{
@@ -811,6 +811,7 @@
"metadata": {},
"outputs": [],
"source": [
+ "%load_ext juliacall\n",
"%julia using Primes: prime"
]
},
@@ -1494,7 +1495,7 @@
"id": "S5dO61g1bDhk"
},
"source": [
- "The full list of PySR parameters can be found here: https://ai.damtp.cam.ac.uk/pysr/api"
+ "The full list of PySR parameters can be found here: https://pysr.ai/api"
]
}
],
diff --git a/examples/search_trace.py b/examples/search_trace.py
new file mode 100644
index 000000000..84228668f
--- /dev/null
+++ b/examples/search_trace.py
@@ -0,0 +1,159 @@
+"""Record the genealogy of a search: every expression evaluated, and its parents.
+
+`use_tracing` makes SymbolicRegression.jl write one JSONL record per (iteration, island)
+holding the live population and every mutation, crossover, tuning and death event. Joining
+those events on member refs gives the full parent-to-child graph, from which the winner's
+ancestry is a walk back up the parent edges.
+
+The trace writer lives in a JSON3 weak-dependency extension, which PySR installs into its
+Julia environment the first time a search asks for it.
+
+Trace records name features `x1`, `x2`, ... in column order, so `x1` below is `X[:, 0]`.
+"""
+
+import collections
+import json
+import pathlib
+import sys
+
+import numpy as np
+
+from pysr import PySRRegressor
+
+rng = np.random.default_rng(0)
+X = rng.uniform(-3, 3, size=(200, 3))
+y = 2.5382 * np.cos(X[:, 2]) + X[:, 0] ** 2 - 1.5
+y = y + rng.normal(0, 0.05 * y.std(), size=200)
+
+TRACE_PATH = "search_trace.jsonl"
+NITERATIONS = 25
+
+SETTINGS = dict(
+ binary_operators=["+", "-", "*", "/"],
+ unary_operators=["cos", "exp"],
+ population_size=27,
+ populations=8,
+ ncycles_per_iteration=90,
+ maxsize=20,
+ random_state=0,
+)
+
+
+def run_search(trace_path, niterations=NITERATIONS):
+ """Run the search with tracing on, leaving the JSONL at `trace_path`."""
+ model = PySRRegressor(
+ niterations=niterations,
+ use_tracing=True,
+ tracing_file=trace_path,
+ **SETTINGS,
+ )
+ model.fit(X, y)
+ return model
+
+
+def load_trace(trace_path):
+ """Read the JSONL into `nodes` (ref -> expression record) and `edges` (parent to child).
+
+ A crossover event is pushed onto both parents' event lists, so it is read from
+ `parent1` only. Each child records which of the two it descends from, so the edge is
+ oriented onto that parent and the other is carried alongside it as `mate`.
+ """
+ nodes, edges = {}, []
+ for line in pathlib.Path(trace_path).read_text().splitlines():
+ record = json.loads(line)
+ if record["record_type"] != "iteration":
+ continue
+ for member in record["members"]:
+ nodes.setdefault(member["ref"], {}).update(member)
+ for ref, entry in record["mutations"].items():
+ parent = int(ref)
+ nodes.setdefault(parent, {}).update(
+ {k: entry[k] for k in ("tree", "loss", "cost", "parent") if k in entry}
+ )
+ for event in entry["events"]:
+ kind = event["type"]
+ detail = event.get("mutation") or event.get("details") or {}
+ if kind == "crossover" and event["parent1"] != parent:
+ continue
+ children = (
+ (event["child1"], event["child2"])
+ if kind == "crossover"
+ else (event["child"],) if "child" in event else ()
+ )
+ for child in children:
+ edges.append(
+ dict(
+ parent=parent,
+ child=child,
+ kind=kind,
+ mutation=detail.get("type"),
+ result=detail.get("result", "accept"),
+ mate=event.get("parent2"),
+ )
+ )
+ for edge in edges:
+ if (
+ edge["mate"] is not None
+ and nodes.get(edge["child"], {}).get("parent") == edge["mate"]
+ ):
+ edge["parent"], edge["mate"] = edge["mate"], edge["parent"]
+ return nodes, edges
+
+
+def ancestry(nodes, edges):
+ """The lit path: the lowest-loss expression walked back up its parent edges."""
+ first_edge = {}
+ for edge in edges:
+ first_edge.setdefault(edge["child"], edge)
+ winner = min(
+ (ref for ref, n in nodes.items() if n.get("loss") is not None),
+ key=lambda ref: nodes[ref]["loss"],
+ )
+ chain, current, seen = [], winner, set()
+ while current in first_edge and current not in seen:
+ seen.add(current)
+ chain.append(first_edge[current])
+ current = first_edge[current]["parent"]
+ chain.reverse()
+ return winner, chain
+
+
+def check(nodes, edges):
+ """A usable genealogy: every edge resolves, the winner chains back to a seed, and
+ crossovers carry a distinct second parent."""
+ if not edges or not any(e["kind"] == "crossover" for e in edges):
+ return False
+ if not any(e["mate"] is not None and e["mate"] != e["parent"] for e in edges):
+ return False
+ if any(e["parent"] not in nodes for e in edges):
+ return False
+ winner, chain = ancestry(nodes, edges)
+ return bool(chain) and nodes[chain[0]["parent"]].get("parent") == -1
+
+
+def main():
+ trace_path = sys.argv[1] if len(sys.argv) > 1 else TRACE_PATH
+ run_search(trace_path)
+ nodes, edges = load_trace(trace_path)
+
+ kinds = collections.Counter(e["kind"] for e in edges)
+ print(f"trace file: {trace_path} ({pathlib.Path(trace_path).stat().st_size} bytes)")
+ print(f"expressions recorded: {len(nodes)}")
+ print(f"parent-to-child edges: {len(edges)} {dict(kinds)}")
+ print(f"rejected: {sum(1 for e in edges if e['result'] == 'reject')}")
+ print(
+ f"edges with a second parent: {sum(1 for e in edges if e['mate'] is not None and e['mate'] != e['parent'])}"
+ )
+
+ winner, chain = ancestry(nodes, edges)
+ print(f"\nwinner loss {nodes[winner]['loss']:.6g}, ancestry {len(chain)} steps")
+ for edge in chain[-8:]:
+ print(
+ f" {edge['kind']:10s} {str(edge['mutation']):32s} -> {nodes[edge['child']].get('tree', '?')[:70]}"
+ )
+
+ print(f"\nusable genealogy: {check(nodes, edges)}")
+
+
+if __name__ == "__main__":
+ main()
diff --git a/examples/test_examples.py b/examples/test_examples.py
new file mode 100644
index 000000000..a4c431c3e
--- /dev/null
+++ b/examples/test_examples.py
@@ -0,0 +1,274 @@
+"""Executable coverage for the shipped documentation examples.
+
+Every example in `examples/` is a plain script a reader can run directly. Most also
+expose the handful of names this suite needs, so the documented result is checked by
+running the documented search rather than by re-describing it here:
+
+ X, y the problem
+ VARIABLE_NAMES list[str], or absent
+ MODEL_KWARGS every PySRRegressor argument except random_state
+ check(model) -> bool the example's own definition of success
+
+Examples with their own runners are registered separately below.
+
+Wall times are single-seed measurements on a `genx` node, one thread, PySR 2.1.0.
+They pick the marker, not an assertion: timing a search on unknown CI hardware would
+be a flaky test.
+"""
+
+import ast
+import importlib.util
+import re
+import sys
+from pathlib import Path
+
+import pytest
+
+EXAMPLES_DIR = Path(__file__).resolve().parent
+
+# name -> slowest measured single-seed wall time in seconds
+EXAMPLES = {
+ "adaptive_mutation_weights": 28,
+ "affine_cipher": 5,
+ "game_of_life": 73,
+ "automatic_batching": 90,
+ "backsolve_mutation": 179,
+ "morton_interleave": 222,
+ "mutations_and_plugins": 214,
+ "any_arity": 362,
+ "turtle_graphics": 2414,
+ "cartpole_objective": 6866,
+ "prng_period": 3230,
+ "magnetic_field": 5507,
+}
+
+# Examples that cannot go through a single PySRRegressor fit and carry their own runner.
+SELF_DRIVEN = {"search_trace", "agent_loop_guesses", "hopper_controller"}
+
+# Examples whose search recovers its target on some seeds and not others. The measured
+# rate is reported in the docs; asserting recovery on one fixed seed would bake in a
+# lucky draw, so these are checked for completing and producing a usable front.
+STOCHASTIC = {"affine_freemod": 190}
+
+FAST_SECONDS = 100
+
+
+def load(name):
+ path = EXAMPLES_DIR / f"{name}.py"
+ spec = importlib.util.spec_from_file_location(name, path)
+ module = importlib.util.module_from_spec(spec)
+ sys.modules[name] = module
+ spec.loader.exec_module(module)
+ return module
+
+
+def run(name):
+ from pysr import PySRRegressor
+
+ module = load(name)
+ model = PySRRegressor(**module.MODEL_KWARGS, random_state=0)
+
+ fit_kwargs = {}
+ if getattr(module, "VARIABLE_NAMES", None):
+ fit_kwargs["variable_names"] = module.VARIABLE_NAMES
+ model.fit(module.X, module.y, **fit_kwargs)
+
+ assert module.check(model), f"{name} no longer reproduces its documented result"
+
+
+@pytest.mark.parametrize("name", [n for n, s in EXAMPLES.items() if s <= FAST_SECONDS])
+def test_example(name):
+ run(name)
+
+
+@pytest.mark.slow
+@pytest.mark.parametrize("name", [n for n, s in EXAMPLES.items() if s > FAST_SECONDS])
+def test_example_slow(name):
+ run(name)
+
+
+@pytest.mark.slow
+def test_agent_loop_reaches_floor():
+ """The loop's premise: a cold search misses the floor and fed-back guesses reach it."""
+ module = load("agent_loop_guesses")
+ rounds = module.main(seed=0)
+
+ assert len(rounds) > 1, "a cold start alone succeeded; the example has no loop left"
+ assert rounds[0]["loss"] > 2 * module.NOISE_FLOOR, "round 0 already at the floor"
+ assert (
+ rounds[-1]["loss"] <= 2 * module.NOISE_FLOOR
+ ), "the loop never reached the floor"
+
+
+@pytest.mark.slow
+def test_search_trace_genealogy(tmp_path):
+ """The trace resolves into a graph the winner's ancestry can be walked back through.
+
+ `use_tracing` pulls JSON3 into the PySR Julia environment on first use, so a cold
+ environment pays an install here before the search starts.
+ """
+ module = load("search_trace")
+ trace_path = str(tmp_path / "search_trace.jsonl")
+
+ module.run_search(trace_path, niterations=3)
+ nodes, edges = module.load_trace(trace_path)
+
+ assert module.check(nodes, edges), "the recorded genealogy no longer resolves"
+
+
+@pytest.mark.slow
+@pytest.mark.parametrize("name", sorted(STOCHASTIC))
+def test_stochastic_example_completes(name):
+ """The search runs and yields a front, whether or not this seed finds the answer."""
+ from pysr import PySRRegressor
+
+ module = load(name)
+ model = PySRRegressor(**module.MODEL_KWARGS, random_state=0)
+ model.fit(module.X, module.y)
+
+ assert not model.equations_.empty, f"{name} produced no front"
+ assert len(module.decode(model)) == len(module.PLAIN)
+
+
+def test_every_example_is_registered():
+ """A new file in examples/ must be registered, or it ships with no coverage."""
+ harness = {Path(__file__).stem, "conftest"}
+ on_disk = {p.stem for p in EXAMPLES_DIR.glob("*.py")} - harness
+ known = set(EXAMPLES) | SELF_DRIVEN | set(STOCHASTIC)
+ assert (
+ on_disk == known
+ ), f"unregistered: {sorted(on_disk - known)}, missing from disk: {sorted(known - on_disk)}"
+
+
+# Examples are ordered by what a reader needs first, not by which part of PySR they
+# exercise: the core of a search, then the knobs that shape one, then value types and
+# the searches built on them. Reordering is a deliberate act, so this list, the index
+# page, and the VitePress sidebar move together.
+READING_ORDER = [
+ "getting-started",
+ "expression-specifications",
+ "objectives",
+ "physics",
+ "search-behaviour",
+ "instrumentation",
+ "value-types",
+ "beyond-numeric-values",
+]
+
+DOCS_SRC = EXAMPLES_DIR.parent / "docs" / "src"
+
+
+def _squash(text):
+ return re.sub(r"[^a-z0-9]", "", text.lower())
+
+
+def test_docs_examples_follow_one_declared_order():
+ """Sidebar, index, and page bodies agree on the reading order, core material first."""
+ if not DOCS_SRC.exists():
+ pytest.skip("docs tree absent: examples were copied out of the repo")
+
+ config = (DOCS_SRC / ".vitepress" / "config.mts").read_text()
+ assert re.findall(r"link: '/examples/([a-z-]+)'", config) == READING_ORDER
+
+ index = (DOCS_SRC / "examples.md").read_text()
+ assert (
+ re.findall(r"^## \[[^]]+\]\(/examples/([a-z-]+)\)$", index, re.M)
+ == READING_ORDER
+ )
+
+ groups = dict(zip(READING_ORDER, index.split("\n## ")[1:]))
+ for slug in READING_ORDER:
+ listed = re.findall(
+ rf"^- \[[^]]+\]\(/examples/{slug}#([a-z0-9-]+)\)", groups[slug], re.M
+ )
+ page = (DOCS_SRC / "examples" / f"{slug}.md").read_text()
+ headings = [h for h in re.findall(r"^## (.+)$", page, re.M) if h != "Preamble"]
+ assert [_squash(a) for a in listed] == [
+ _squash(h) for h in headings
+ ], f"{slug}: index lists {listed}, page has {headings}"
+
+
+def _split_literals(source):
+ """Literals written as adjacent parts across lines, which hide their own seams."""
+ found = []
+ for node in ast.walk(ast.parse(source)):
+ if isinstance(node, ast.JoinedStr) or (
+ isinstance(node, ast.Constant) and isinstance(node.value, str)
+ ):
+ if node.lineno == node.end_lineno:
+ continue
+ segment = ast.get_source_segment(source, node) or ""
+ if not re.match(r"^[fFrRbB]{0,2}(\"\"\"|''')", segment):
+ found.append(segment.splitlines()[0].strip())
+ return found
+
+
+def test_no_string_is_split_across_lines():
+ """A string spanning lines is triple-quoted, never adjacent parts.
+
+ Implicit concatenation hides whitespace at the seams. These files are Julia source
+ a reader copies, where a dropped space silently changes the program.
+ """
+ offenders = {}
+ for path in sorted(EXAMPLES_DIR.glob("*.py")):
+ split = _split_literals(path.read_text())
+ if split:
+ offenders[path.name] = split
+ for page in sorted((DOCS_SRC / "examples").glob("*.md")):
+ for block in re.findall(
+ r"^```python\n(.*?)^```", page.read_text(), re.S | re.M
+ ):
+ try:
+ split = _split_literals(block)
+ except SyntaxError:
+ continue
+ if split:
+ offenders.setdefault(page.name, []).extend(split)
+ assert not offenders, f"strings split across lines: {offenders}"
+
+
+def _stray_indents(source):
+ """Multi-line literals whose body resets to a shallower column than its own opener."""
+ lines = source.split("\n")
+ found = []
+ for node in ast.walk(ast.parse(source)):
+ if not isinstance(node, (ast.JoinedStr, ast.Constant)):
+ continue
+ if isinstance(node, ast.Constant) and not isinstance(node.value, str):
+ continue
+ if node.lineno == node.end_lineno:
+ continue
+ opener = lines[node.lineno - 1]
+ body = [line for line in lines[node.lineno : node.end_lineno] if line.strip()]
+ if not body:
+ continue
+ base = len(opener) - len(opener.lstrip(" "))
+ shallowest = min(len(line) - len(line.lstrip(" ")) for line in body)
+ if shallowest != base:
+ found.append(opener.strip())
+ return found
+
+
+def test_embedded_julia_is_laid_out_like_its_host():
+ """A Julia block inside a Python file is indented with the Python around it.
+
+ Julia ignores indentation, so an `end` dragged out to column 0 buys nothing and
+ leaves the Python file looking broken at every closer.
+ """
+ offenders = {}
+ for path in sorted(EXAMPLES_DIR.glob("*.py")):
+ stray = _stray_indents(path.read_text())
+ if stray:
+ offenders[path.name] = stray
+ if DOCS_SRC.exists():
+ for page in sorted((DOCS_SRC / "examples").glob("*.md")):
+ for block in re.findall(
+ r"^```python\n(.*?)^```", page.read_text(), re.S | re.M
+ ):
+ try:
+ stray = _stray_indents(block)
+ except SyntaxError:
+ continue
+ if stray:
+ offenders.setdefault(page.name, []).extend(stray)
+ assert not offenders, f"literals dedented below their opener: {offenders}"
diff --git a/examples/turtle_graphics.py b/examples/turtle_graphics.py
new file mode 100644
index 000000000..7335100bf
--- /dev/null
+++ b/examples/turtle_graphics.py
@@ -0,0 +1,355 @@
+"""Search for a turtle program that draws a plus-sign outline.
+
+The value the search carries is a drawing: a list of pen commands, `forward` and
+`turn`. The operators move and turn the pen, so every expression is a program and
+its value is the picture that program draws. The loss rasterizes both drawings
+and compares them with a chamfer distance.
+"""
+
+import math
+
+import numpy as np
+
+from pysr import PySRRegressor, TypeSpec
+
+# --------------------------------------------------------------- the turtle
+# A Path is a flat command list carried as two parallel vectors: `kind` (0
+# forward, 1 turn) and `val` (a length, or an angle in radians).
+PREAMBLE = r"""
+import Random
+
+const _GRID = 96
+const _MARGIN = 3
+# Off-centre by an arbitrary fraction of a pixel: axis-aligned figures otherwise
+# land exactly on rounding ties, where a one-ulp change of an angle flips a pixel
+# and the loss stops being reproducible across implementations.
+const _OFFSET = 0.1373
+const _MAXCMD = 8192
+const _FAIL = 4.0
+const _DTCACHE = Dict{UInt64,Tuple{BitMatrix,Matrix{Float64}}}()
+const _DTLOCK = ReentrantLock()
+
+_empty() = (Int8[], Float64[])
+
+# ----------------------------------------------------- constants and printing
+# The width of the mutation step at temperature 1, in degrees.
+const _SIGMA_DEG = 30.0
+
+# One turn, any real angle in (-180, 180].
+_real(rng) = deg2rad(180.0 - 360.0 * rand(rng))
+
+function _wrap(deg)
+ d = mod(deg + 180.0, 360.0) - 180.0
+ return d <= -180.0 ? 180.0 : d
+end
+
+_perturb(rng, angle, temperature) =
+ deg2rad(_wrap(rad2deg(angle) + _SIGMA_DEG * temperature * Random.randn(rng)))
+
+# The optimizer sees exactly the turns: one Float64 per turn, forwards untouched.
+_turns(kind::Vector{Int8}, val::Vector{Float64}) =
+ [val[i] for i in eachindex(kind) if kind[i] == Int8(1)]
+
+function _put_turns(kind::Vector{Int8}, val::Vector{Float64}, c)
+ out = copy(val)
+ j = 0
+ for i in eachindex(kind)
+ if kind[i] == Int8(1)
+ j += 1
+ out[i] = c[j]
+ end
+ end
+ return out
+end
+
+function _degstr(angle)
+ d = rad2deg(angle)
+ r = round(d)
+ return abs(d - r) < 1e-9 ? string(Int(r)) : string(d)
+end
+
+_cmd(k::Int8, v::Float64) =
+ k == Int8(1) ? "t" * _degstr(v) : (v == 1.0 ? "f" : "f" * string(v))
+
+# Constant folding turns an all-constant subtree into one multi-command Path, so
+# a constant has to print as something that can be read back command by command.
+function _literal(kind::Vector{Int8}, val::Vector{Float64})
+ length(kind) == 1 && return _cmd(kind[1], val[1])
+ return "p[" * join([_cmd(kind[i], val[i]) for i in eachindex(kind)], ",") * "]"
+end
+
+# --------------------------------------------------------- the path algebra
+function _endpoint(kind::Vector{Int8}, val::Vector{Float64})
+ x = 0.0; y = 0.0; c = 1.0; s = 0.0
+ @inbounds for i in eachindex(kind)
+ if kind[i] == Int8(0)
+ x += c * val[i]; y += s * val[i]
+ else
+ ca = cos(val[i]); sa = sin(val[i])
+ c, s = c * ca - s * sa, s * ca + c * sa
+ end
+ end
+ return x, y
+end
+
+function _seq(ak::Vector{Int8}, av::Vector{Float64},
+ bk::Vector{Int8}, bv::Vector{Float64})
+ (isempty(ak) && isempty(bk)) && return _empty()
+ length(ak) + length(bk) > _MAXCMD && return _empty()
+ return (vcat(ak, bk), vcat(av, bv))
+end
+
+function _mir(kind::Vector{Int8}, val::Vector{Float64})
+ isempty(kind) && return _empty()
+ out = copy(val)
+ @inbounds for i in eachindex(kind)
+ kind[i] == Int8(1) && (out[i] = -out[i])
+ end
+ return (copy(kind), out)
+end
+
+# L-system self-application: replace every forward by a copy of the whole path,
+# scaled so the copy's net displacement equals that forward.
+function _nest(kind::Vector{Int8}, val::Vector{Float64})
+ isempty(kind) && return _empty()
+ nf = count(==(Int8(0)), kind)
+ nf == 0 && return _empty()
+ n = length(kind)
+ n + nf * (n - 1) > _MAXCMD && return _empty()
+ x, y = _endpoint(kind, val)
+ d = hypot(x, y)
+ d < 1e-9 && return _empty()
+ ok = Vector{Int8}(); ov = Vector{Float64}()
+ sizehint!(ok, n + nf * (n - 1)); sizehint!(ov, n + nf * (n - 1))
+ @inbounds for i in eachindex(kind)
+ if kind[i] == Int8(0)
+ r = val[i] / d
+ for j in eachindex(kind)
+ push!(ok, kind[j])
+ push!(ov, kind[j] == Int8(0) ? val[j] * r : val[j])
+ end
+ else
+ push!(ok, kind[i]); push!(ov, val[i])
+ end
+ end
+ return (ok, ov)
+end
+
+# ------------------------------------------------------------------ the loss
+# Walk the commands and stamp every drawn segment onto a 96x96 grid, after
+# centring the bounding box and scaling its larger extent to the grid, so the
+# comparison sees shape and not place or size.
+function _raster(kind::Vector{Int8}, val::Vector{Float64})
+ isempty(kind) && return nothing
+ nf = count(==(Int8(0)), kind)
+ nf == 0 && return nothing
+ xs = Vector{Float64}(undef, 4 * nf)
+ x = 0.0; y = 0.0; c = 1.0; s = 0.0
+ j = 0
+ xmin = Inf; xmax = -Inf; ymin = Inf; ymax = -Inf
+ @inbounds for i in eachindex(kind)
+ if kind[i] == Int8(0)
+ x2 = x + c * val[i]; y2 = y + s * val[i]
+ xs[4j + 1] = x; xs[4j + 2] = y; xs[4j + 3] = x2; xs[4j + 4] = y2
+ j += 1
+ xmin = min(xmin, x, x2); xmax = max(xmax, x, x2)
+ ymin = min(ymin, y, y2); ymax = max(ymax, y, y2)
+ x = x2; y = y2
+ else
+ ca = cos(val[i]); sa = sin(val[i])
+ c, s = c * ca - s * sa, s * ca + c * sa
+ end
+ end
+ (isfinite(xmin) && isfinite(xmax) && isfinite(ymin) && isfinite(ymax)) ||
+ return nothing
+ scale = max(xmax - xmin, ymax - ymin)
+ (scale < 1e-9 || !isfinite(scale)) && return nothing
+ span = _GRID - 1 - 2 * _MARGIN
+ cx = (xmin + xmax) / 2; cy = (ymin + ymax) / 2
+ mask = falses(_GRID, _GRID)
+ @inbounds for k in 0:(nf - 1)
+ u0 = _MARGIN + _OFFSET + ((xs[4k + 1] - cx) / scale + 0.5) * span
+ v0 = _MARGIN + _OFFSET + ((xs[4k + 2] - cy) / scale + 0.5) * span
+ u1 = _MARGIN + _OFFSET + ((xs[4k + 3] - cx) / scale + 0.5) * span
+ v1 = _MARGIN + _OFFSET + ((xs[4k + 4] - cy) / scale + 0.5) * span
+ steps = max(1, ceil(Int, 2 * hypot(u1 - u0, v1 - v0)))
+ for t in 0:steps
+ a = t / steps
+ i = clamp(round(Int, u0 + a * (u1 - u0)) + 1, 1, _GRID)
+ jj = clamp(round(Int, v0 + a * (v1 - v0)) + 1, 1, _GRID)
+ mask[i, jj] = true
+ end
+ end
+ return mask
+end
+
+# Exact squared Euclidean distance transform, one dimension at a time.
+function _dt1d!(d::Vector{Float64}, f::Vector{Float64}, n::Int,
+ v::Vector{Int}, z::Vector{Float64})
+ k = 1; v[1] = 1; z[1] = -1e20; z[2] = 1e20
+ @inbounds for q in 2:n
+ s = ((f[q] + q * q) - (f[v[k]] + v[k] * v[k])) / (2q - 2v[k])
+ while s <= z[k]
+ k -= 1
+ s = ((f[q] + q * q) - (f[v[k]] + v[k] * v[k])) / (2q - 2v[k])
+ end
+ k += 1; v[k] = q; z[k] = s; z[k + 1] = 1e20
+ end
+ k = 1
+ @inbounds for q in 1:n
+ while z[k + 1] < q
+ k += 1
+ end
+ d[q] = (q - v[k])^2 + f[v[k]]
+ end
+ return d
+end
+
+function _edt(mask::BitMatrix)
+ n1, n2 = size(mask)
+ big = 1e12
+ d2 = Matrix{Float64}(undef, n1, n2)
+ @inbounds for j in 1:n2, i in 1:n1
+ d2[i, j] = mask[i, j] ? 0.0 : big
+ end
+ m = max(n1, n2)
+ f = Vector{Float64}(undef, m); d = Vector{Float64}(undef, m)
+ v = Vector{Int}(undef, m + 1); z = Vector{Float64}(undef, m + 2)
+ @inbounds for j in 1:n2
+ for i in 1:n1; f[i] = d2[i, j]; end
+ _dt1d!(d, f, n1, v, z)
+ for i in 1:n1; d2[i, j] = d[i]; end
+ end
+ @inbounds for i in 1:n1
+ for j in 1:n2; f[j] = d2[i, j]; end
+ _dt1d!(d, f, n2, v, z)
+ for j in 1:n2; d2[i, j] = d[j]; end
+ end
+ return d2
+end
+
+# The target never changes, so its raster and its transform are computed once.
+function _target_grids(kind::Vector{Int8}, val::Vector{Float64})
+ key = hash(kind, hash(val))
+ lock(_DTLOCK) do
+ get!(_DTCACHE, key) do
+ mask = _raster(kind, val)
+ mask === nothing ? (falses(_GRID, _GRID), fill(1e12, _GRID, _GRID)) :
+ (mask, _edt(mask))
+ end
+ end
+end
+
+function _mean_to(mask::BitMatrix, d2::Matrix{Float64})
+ total = 0.0; n = 0
+ @inbounds for i in eachindex(mask)
+ if mask[i]
+ total += sqrt(d2[i]); n += 1
+ end
+ end
+ return n == 0 ? nothing : total / n
+end
+
+# Symmetric chamfer distance in pixels, divided by the grid width. Zero means
+# the two drawings lit exactly the same pixels.
+function _chamfer(pk::Vector{Int8}, pv::Vector{Float64},
+ tk::Vector{Int8}, tv::Vector{Float64})
+ a = _raster(pk, pv)
+ a === nothing && return _FAIL
+ b, db = _target_grids(tk, tv)
+ any(b) || return _FAIL
+ da = _edt(a)
+ ab = _mean_to(a, db); ba = _mean_to(b, da)
+ (ab === nothing || ba === nothing) && return _FAIL
+ v = 0.5 * (ab + ba) / _GRID
+ return isfinite(v) ? v : _FAIL
+end
+"""
+
+PATH = TypeSpec(
+ "Path",
+ fields={"kind": "Vector{Int8}", "val": "Vector{Float64}"},
+ preamble=PREAMBLE,
+ # A constant is one turn, any real angle in (-180, 180].
+ sample="rng -> Path(Int8[1], [_real(rng)])",
+ # Perturb one of the value's angles by a temperature-scaled step, or
+ # resample it outright. Constant folding can hand this hook a value of
+ # several commands, so it moves a command in place.
+ mutate="""(rng, value, temperature) -> begin
+ turns = findall(isone, value.kind)
+ isempty(turns) && return Path(Int8[1], [_real(rng)])
+ i = rand(rng, turns)
+ val = copy(value.val)
+ val[i] = rand(rng, Bool) ? _perturb(rng, val[i], temperature) : _real(rng)
+ Path(copy(value.kind), val)
+ end""",
+ # Every turn in the value is a scalar the constant optimizer may fit.
+ scalar_constants="value -> _turns(value.kind, value.val)",
+ with_scalar_constants="(value, c) -> Path(copy(value.kind), _put_turns(value.kind, value.val, c))",
+ is_valid="value -> !isempty(value.kind) && all(isfinite, value.val)",
+ string="value -> _literal(value.kind, value.val)",
+)
+
+OPERATORS = {
+ 1: [
+ "dup(a::Path) = (t = _seq(a.kind, a.val, a.kind, a.val); Path(t[1], t[2]))",
+ "tri(a::Path) = (t = _seq(a.kind, a.val, a.kind, a.val); u = _seq(t[1], t[2], a.kind, a.val); Path(u[1], u[2]))",
+ "mir(a::Path) = (t = _mir(a.kind, a.val); Path(t[1], t[2]))",
+ "nest(a::Path) = (t = _nest(a.kind, a.val); Path(t[1], t[2]))",
+ ],
+ 2: [
+ "seq(a::Path, b::Path) = (t = _seq(a.kind, a.val, b.kind, b.val); Path(t[1], t[2]))"
+ ],
+}
+
+CHAMFER_LOSS = "chamfer(prediction::Path, target::Path)::Float64 = _chamfer(prediction.kind, prediction.val, target.kind, target.val)"
+
+
+def path(commands):
+ """A ('F', length) / ('T', degrees) command list as the two Path vectors."""
+ kind = [0 if op == "F" else 1 for op, _ in commands]
+ val = [a if op == "F" else math.radians(a) for op, a in commands]
+ return kind, val
+
+
+# The target: a plus-sign outline, four copies of a three-edge arm.
+ARM = [("F", 1.0), ("T", 90.0), ("F", 1.0), ("T", -90.0), ("F", 1.0), ("T", 90.0)]
+TARGET = ARM * 4
+
+# One row. The only feature is `f`, a single unit step forward.
+X = np.empty((1, 1), dtype=object)
+X[0, 0] = ([0], [1.0])
+y = np.empty(1, dtype=object)
+y[0] = path(TARGET)
+VARIABLE_NAMES = ["f"]
+
+MODEL_KWARGS = dict(
+ type_spec=PATH,
+ operators=OPERATORS,
+ elementwise_loss=CHAMFER_LOSS,
+ niterations=20,
+ verbosity=0,
+)
+
+EXACT = 1e-9
+
+
+def check(model):
+ """Some program on the front draws the target, to floating-point precision."""
+ return float(model.equations_["loss"].min()) < EXACT
+
+
+def main():
+ model = PySRRegressor(**MODEL_KWARGS, random_state=0)
+ model.fit(X, y, variable_names=VARIABLE_NAMES)
+ print(model.equations_[["complexity", "loss", "equation"]].to_string(index=False))
+ exact = model.equations_[model.equations_["loss"] < EXACT]
+ if len(exact):
+ best = exact.loc[exact["complexity"].idxmin()]
+ print(f"\ndraws the target: {best.equation}")
+ else:
+ print("\nno program drew the target")
+
+
+if __name__ == "__main__":
+ main()
diff --git a/pyproject.toml b/pyproject.toml
index fccd37f7a..a19729970 100644
--- a/pyproject.toml
+++ b/pyproject.toml
@@ -4,25 +4,25 @@ build-backend = "hatchling.build"
[project]
name = "pysr"
-version = "2.0.0a2"
+version = "2.6.0"
authors = [
{name = "Miles Cranmer", email = "miles.cranmer@gmail.com"},
]
description = "Simple and efficient symbolic regression"
readme = {file = "README.md", content-type = "text/markdown"}
license = {file = "LICENSE"}
-requires-python = ">=3.9"
+requires-python = ">=3.10"
classifiers = [
"Programming Language :: Python :: 3",
"Operating System :: OS Independent",
"License :: OSI Approved :: Apache Software License"
]
dependencies = [
- "sympy>=1.0.0,<2.0.0",
+ "sympy>=1.14.0,<2.0.0",
"pandas>=0.21.0,<4.0.0",
"numpy>=1.13.0,<3.0.0",
"scikit_learn>=1.0.0,<2.0.0",
- "juliacall>=0.9.28,<0.9.29",
+ "juliacall>=0.9.36,<0.9.37",
"click>=7.0.0,<9.0.0",
"typing-extensions>=4.0.0,<5.0.0",
]
@@ -35,18 +35,18 @@ docs = [
dev = [
"coverage>=7,<8",
"beartype>=0.19,<0.23",
- "ipykernel>=6,<7",
- "ipython>=8,<9",
- "jax[cpu]>=0.4,<0.6",
+ "ipykernel>=6,<8",
+ "ipython>=8,<10",
+ "jax[cpu]>=0.4,<0.12",
"jupyter>=1,<2",
"mypy>=1,<3",
"nbval>=0.11,<0.12",
"pandas-stubs",
"pre-commit>=3.0,<5",
- "pytest-cov>=5,<7",
+ "pytest-cov>=5,<8",
"pytest>=8,<10",
"tensorboard>=2,<3",
- "torch>=2,<3",
+ "torch>=2.7,<3",
"types-openpyxl",
"types-pytz",
]
diff --git a/pysr/__init__.py b/pysr/__init__.py
index 3e3f7c8be..8a1f5cb36 100644
--- a/pysr/__init__.py
+++ b/pysr/__init__.py
@@ -29,12 +29,35 @@
from .expression_specs import (
AbstractExpressionSpec,
ExpressionSpec,
- ParametricExpressionSpec,
TemplateExpressionSpec,
)
from .julia_extensions import load_all_packages
from .logger_specs import AbstractLoggerSpec, TensorBoardLoggerSpec
+from .mutations import (
+ AbstractMutation,
+ AddNodeMutation,
+ BacksolveMutation,
+ ConstantMutation,
+ DeleteNodeMutation,
+ DoNothingMutation,
+ FeatureMutation,
+ InsertNodeMutation,
+ OperatorMutation,
+ OptimizeMutation,
+ RandomizeMutation,
+ RotateTreeMutation,
+ SimplifyMutation,
+ SwapOperandsMutation,
+)
+from .plugins import (
+ AbstractPlugin,
+ AdaptiveMutationWeightsPlugin,
+ AdaptiveParsimonyPlugin,
+ MutationBurstPlugin,
+ SimulatedAnnealingPlugin,
+)
from .sr import PySRRegressor
+from .type_specs import TypeSpec
try:
__version__ = version("pysr")
@@ -54,9 +77,28 @@
"AbstractExpressionSpec",
"ExpressionSpec",
"TemplateExpressionSpec",
- "ParametricExpressionSpec",
"AbstractLoggerSpec",
"TensorBoardLoggerSpec",
+ "TypeSpec",
+ "AbstractMutation",
+ "ConstantMutation",
+ "OperatorMutation",
+ "FeatureMutation",
+ "SwapOperandsMutation",
+ "AddNodeMutation",
+ "InsertNodeMutation",
+ "DeleteNodeMutation",
+ "RotateTreeMutation",
+ "BacksolveMutation",
+ "SimplifyMutation",
+ "RandomizeMutation",
+ "OptimizeMutation",
+ "DoNothingMutation",
+ "AbstractPlugin",
+ "AdaptiveMutationWeightsPlugin",
+ "AdaptiveParsimonyPlugin",
+ "MutationBurstPlugin",
+ "SimulatedAnnealingPlugin",
"best",
"best_callable",
"best_row",
diff --git a/pysr/_cli/main.py b/pysr/_cli/main.py
index cd81710ff..f1a06fbca 100644
--- a/pysr/_cli/main.py
+++ b/pysr/_cli/main.py
@@ -1,4 +1,5 @@
import fnmatch
+import os
import sys
import unittest
import warnings
@@ -10,6 +11,7 @@
runtests,
runtests_autodiff,
runtests_dev,
+ runtests_interrupt,
runtests_jax,
runtests_slurm,
runtests_startup,
@@ -50,7 +52,17 @@ def _install(julia_project, quiet, precompile):
)
-TEST_OPTIONS = {"main", "jax", "torch", "autodiff", "cli", "dev", "startup", "slurm"}
+TEST_OPTIONS = {
+ "main",
+ "jax",
+ "torch",
+ "autodiff",
+ "cli",
+ "dev",
+ "startup",
+ "slurm",
+ "interrupt",
+}
@pysr.command("test")
@@ -65,9 +77,23 @@ def _install(julia_project, quiet, precompile):
def _tests(tests, expressions):
"""Run parts of the PySR test suite.
- Choose from main, jax, torch, autodiff, cli, dev, startup, and slurm.
+ Choose from main, jax, torch, autodiff, cli, dev, startup, slurm, and interrupt.
You can give multiple tests, separated by commas.
"""
+ try:
+ shard_count = int(os.environ.get("PYSR_TEST_SHARD_COUNT", "1"))
+ shard_index = int(os.environ.get("PYSR_TEST_SHARD_INDEX", "0"))
+ except ValueError as error:
+ raise click.ClickException(
+ "PYSR_TEST_SHARD_COUNT and PYSR_TEST_SHARD_INDEX must be integers."
+ ) from error
+ if shard_count <= 0:
+ raise click.ClickException("PYSR_TEST_SHARD_COUNT must be a positive integer.")
+ if not 0 <= shard_index < shard_count:
+ raise click.ClickException(
+ "PYSR_TEST_SHARD_INDEX must satisfy 0 <= PYSR_TEST_SHARD_INDEX "
+ "< PYSR_TEST_SHARD_COUNT."
+ )
test_cases = []
for test in tests.split(","):
if test == "main":
@@ -85,6 +111,8 @@ def _tests(tests, expressions):
test_cases.extend(runtests_dev(just_tests=True))
elif test == "startup":
test_cases.extend(runtests_startup(just_tests=True))
+ elif test == "interrupt":
+ test_cases.extend(runtests_interrupt(just_tests=True))
elif test == "slurm":
test_cases.extend(runtests_slurm(just_tests=True))
else:
@@ -100,6 +128,7 @@ def _tests(tests, expressions):
for expression in expressions
):
suite.addTest(test)
+ suite = unittest.TestSuite(list(suite)[shard_index::shard_count])
runner = unittest.TextTestRunner()
results = runner.run(suite)
diff --git a/pysr/expression_specs.py b/pysr/expression_specs.py
index bd8d5ec71..187516d5d 100644
--- a/pysr/expression_specs.py
+++ b/pysr/expression_specs.py
@@ -1,11 +1,10 @@
from __future__ import annotations
import copy
-import textwrap
-import warnings
from abc import ABC, abstractmethod
+from dataclasses import dataclass
from textwrap import dedent
-from typing import TYPE_CHECKING, Any, NewType, overload
+from typing import TYPE_CHECKING, Any, ClassVar, NewType
import numpy as np
import pandas as pd
@@ -13,7 +12,6 @@
from .export import add_export_formats
from .julia_helpers import jl_array
from .julia_import import AnyValue, SymbolicRegression, jl
-from .utils import ArrayLike
try:
from typing import TypeAlias
@@ -42,9 +40,7 @@ class AbstractExpressionSpec(ABC):
2. create_exports(), which will be used to create the exports of the equations, such as
the executable format, the SymPy format, etc.
- It may also optionally implement:
-
- - supports_sympy, supports_torch, supports_jax, supports_latex: Whether this expression type supports the corresponding export format.
+ Implementations may also declare the export formats they support.
"""
@abstractmethod
@@ -52,6 +48,30 @@ def julia_expression_spec(self) -> AnyValue:
"""The expression specification"""
pass # pragma: no cover
+ def _julia_expression_spec_source(self, *, prototype: str | None) -> str | None:
+ """Return self-contained Julia source for a TypeSpec-compatible spec.
+
+ ``prototype`` is Julia source for one value of the generated type. PySR
+ evaluates the returned source once in the fingerprinted runtime module.
+ """
+ return None
+
+ def _julia_expression_spec_function_selector(self) -> str | None:
+ """Return Julia source selecting the callable nested in the spec.
+
+ PySR evaluates the source to a selector, invokes it with the expression
+ spec, and binds the result under a deterministic runtime-module name so
+ checkpoints and workers resolve the same callable identity.
+ """
+ return None
+
+ @property
+ def supports_type_spec(self) -> bool:
+ return False
+
+ def _validate_type_spec(self) -> None:
+ pass
+
@abstractmethod
def create_exports(
self,
@@ -108,6 +128,10 @@ def create_exports(
output_torch_format=model.output_torch_format,
)
+ @property
+ def supports_type_spec(self) -> bool:
+ return True
+
@property
def supports_sympy(self):
return True
@@ -125,6 +149,7 @@ def supports_latex(self):
return True
+@dataclass
class TemplateExpressionSpec(AbstractExpressionSpec):
"""Spec for templated expressions.
@@ -177,211 +202,54 @@ class TemplateExpressionSpec(AbstractExpressionSpec):
the derivative of f with respect to its first argument, evaluated at x.
"""
- _spec_cache: dict[tuple[str, ...], AnyValue] = {}
+ combine: str
+ expressions: list[str]
+ variable_names: list[str]
+ parameters: dict[str, int] | None = None
- @overload
- def __init__(
- self,
- function_symbols: list[str],
- combine: str,
- num_features: dict[str, int] | None = None,
- ) -> None: ...
-
- @overload
- def __init__(
- self,
- combine: str,
- *,
- expressions: list[str],
- variable_names: list[str],
- parameters: dict[str, int] | None = None,
- ) -> None: ...
-
- def __init__(
- self,
- *args,
- **kwargs,
- ):
- """Handle both formats with combine as explicit parameter"""
- self._old_format = len(args) >= 2 or "function_symbols" in kwargs
-
- if self._old_format:
- self._load_old_format(*args, **kwargs)
- else:
- self._load_new_format(*args, **kwargs)
-
- def _load_old_format(
- self,
- function_symbols: list[str],
- combine: str,
- num_features: dict[str, int] | None = None,
- ):
- self.function_symbols = function_symbols
- self.combine = combine
- self.num_features = num_features
- # TODO: warn about old format after some versions
-
- def _load_new_format(
- self,
- combine: str,
- *,
- expressions: list[str],
- variable_names: list[str],
- parameters: dict[str, int] | None = None,
- ):
- self.combine = combine
- self.expressions = expressions
- self.variable_names = variable_names
- self.parameters = parameters
+ _spec_cache: ClassVar[dict[tuple[str, ...], AnyValue]] = {}
def _get_cache_key(self):
- if self._old_format:
- return (
- "old",
- str(self.function_symbols),
- self.combine,
- str(self.num_features),
- )
- else:
- return (
- "new",
- self.combine,
- str(self.expressions),
- str(self.variable_names),
- str(self.parameters),
- )
+ return (
+ self.combine,
+ str(self.expressions),
+ str(self.variable_names),
+ str(self.parameters),
+ )
def julia_expression_spec(self):
key = self._get_cache_key()
- if key in self._spec_cache:
- return self._spec_cache[key]
-
- if self._old_format:
- result = SymbolicRegression.TemplateExpressionSpec(
- structure=self.julia_expression_options().structure
- )
- else:
- result = self._call_template_macro()
-
- self._spec_cache[key] = result
- return result
+ if key not in self._spec_cache:
+ self._spec_cache[key] = self._call_template_macro()
+ return self._spec_cache[key]
def _call_template_macro(self):
return jl.seval(self._template_macro_str())
- def _template_macro_str(self):
+ def _template_macro_str(self, *, prototype: str | None = None) -> str:
template_inputs = [f"expressions=({', '.join(self.expressions) + ','})"]
if self.parameters:
template_inputs.append(
f"parameters=({', '.join([f'{p}={self.parameters[p]}' for p in self.parameters]) + ','})"
)
+ if prototype is not None:
+ template_inputs.append(f"prototype={prototype}")
return dedent(f"""
@template_spec({', '.join(template_inputs) + ','}) do {', '.join(self.variable_names)}
{self.combine}
end
""")
- def julia_expression_options(self):
- f_combine = jl.seval(self.combine)
- creator = jl.seval("""
- function _pysr_create_template_structure(
- @nospecialize(function_symbols::AbstractVector),
- @nospecialize(combine::Function),
- @nospecialize(num_features::Union{Nothing,AbstractDict})
- )
- tuple_symbol = (map(Symbol, function_symbols)..., )
- num_features = if num_features === nothing
- nothing
- else
- NamedTuple(Symbol(k) => v for (k, v) in num_features)
- end
- structure = SymbolicRegression.TemplateStructure{tuple_symbol}(combine, num_features)
- return (; structure)
- end
- """)
- return creator(self.function_symbols, f_combine, self.num_features)
+ def _julia_expression_spec_source(self, *, prototype: str | None) -> str:
+ return self._template_macro_str(prototype=prototype)
+
+ def _julia_expression_spec_function_selector(self) -> str:
+ return "spec -> spec.structure.combine"
@property
- def evaluates_in_julia(self):
+ def supports_type_spec(self) -> bool:
return True
- def create_exports(
- self,
- model: PySRRegressor,
- equations: pd.DataFrame,
- search_output,
- i: int | None = None,
- ) -> pd.DataFrame:
- # We try to load the raw julia state from a saved binary stream
- # if not provided.
- search_output = search_output or model.julia_state_
- return _search_output_to_callable_expressions(equations, search_output, i)
-
-
-def parametric_expression_deprecation_warning(
- max_parameters: int, variable_names: ArrayLike[str]
-):
- function_name = "f"
- var_names = list(variable_names)
- message = dedent(f"""
- ParametricExpressionSpec is deprecated. you should switch to TemplateExpressionSpec with explicit parameters indexed by category.
-
- Since you have `max_parameters={max_parameters}` and `variable_names=[{", ".join(f'"{v}"' for v in var_names)}]`, you could migrate like this:
-
- n_categories = len(np.unique(category)) # count the number of parameters required
- expression_spec = TemplateExpressionSpec(
- expressions=["{function_name}"],
- variable_names=[{", ".join(f'"{v}"' for v in var_names + ["category"])}],
- parameters={{{", ".join(f'"p{i+1}": n_categories' for i in range(max_parameters))}}},
- combine="{function_name}({', '.join(var_names + [f'p{i+1}[category]' for i in range(max_parameters)])})",
- )
- X = np.column_stack([X, category]) # add the category column
-
- Finally, do not pass `category` when calling .fit().
- """).strip()
- wrapped = "\n".join(textwrap.fill(line, 88) for line in message.splitlines())
- warnings.warn(wrapped, FutureWarning, stacklevel=3)
-
-
-class ParametricExpressionSpec(AbstractExpressionSpec):
- """Spec for parametric expressions that vary by category.
-
- **This is deprecated in favor of the `TemplateExpressionSpec` class,
- which now supports parameters indexed by category.**
-
- This class allows you to specify expressions with parameters that vary across different
- categories in your dataset. The expression structure remains the same, but parameters
- are optimized separately for each category.
-
- Parameters
- ----------
- max_parameters : int
- Maximum number of parameters that can appear in the expression. Each parameter
- will take on different values for each category in the data.
-
- Examples
- --------
- For example, if we want to allow for a model with up to 2 parameters (each category
- can have a different value for these parameters), we can use:
-
- ```python
- model = PySRRegressor(
- expression_spec=ParametricExpressionSpec(max_parameters=2),
- binary_operators=["+", "*"],
- unary_operators=["sin"]
- )
- model.fit(X, y, category=category)
- ```
- """
-
- def __init__(self, max_parameters: int):
- self.max_parameters = max_parameters
-
- def julia_expression_spec(self):
- return SymbolicRegression.ParametricExpressionSpec(
- max_parameters=self.max_parameters, warn=False
- )
-
@property
def evaluates_in_julia(self):
return True
@@ -392,7 +260,7 @@ def create_exports(
equations: pd.DataFrame,
search_output,
i: int | None = None,
- ):
+ ) -> pd.DataFrame:
search_output = search_output or model.julia_state_
return _search_output_to_callable_expressions(equations, search_output, i)
@@ -407,7 +275,9 @@ def __call__(self, X: np.ndarray, *args):
def _search_output_to_callable_expressions(
- equations: pd.DataFrame, search_output, i: int | None
+ equations: pd.DataFrame,
+ search_output,
+ i: int | None,
) -> pd.DataFrame:
equations = copy.deepcopy(equations)
_, all_out_hof = search_output
diff --git a/pysr/interrupt.py b/pysr/interrupt.py
new file mode 100644
index 000000000..cd1d9b4ad
--- /dev/null
+++ b/pysr/interrupt.py
@@ -0,0 +1,132 @@
+"""Cooperative SIGINT handling so searches stop gracefully instead of killing the process."""
+
+from __future__ import annotations
+
+import _thread
+import ctypes
+import os
+import signal
+import socket
+import threading
+import warnings
+from contextlib import ExitStack, contextmanager
+from ctypes import wintypes
+
+from .julia_import import SymbolicRegression
+
+
+class _SigactionStorage(ctypes.Union):
+ _fields_ = [
+ ("alignment", ctypes.c_longdouble),
+ ("storage", ctypes.c_ubyte * 1024),
+ ]
+
+
+def _libc_with_sigaction():
+ libc = ctypes.CDLL(None, use_errno=True)
+ libc.sigaction.argtypes = [ctypes.c_int, ctypes.c_void_p, ctypes.c_void_p]
+ libc.sigaction.restype = ctypes.c_int
+ return libc
+
+
+def _checked_sigaction(libc, action, old_action):
+ result = libc.sigaction(signal.SIGINT, action, old_action)
+ if result != 0:
+ errno = ctypes.get_errno()
+ raise OSError(errno, os.strerror(errno))
+
+
+@contextmanager
+def _console_ctrl_c_to_python():
+ """Hand console Ctrl-C to Python before Julia's handler exits the process.
+
+ Windows calls console handlers last-registered first, so this one runs first.
+ """
+ kernel32 = ctypes.WinDLL("kernel32", use_last_error=True)
+ HANDLER = ctypes.WINFUNCTYPE(wintypes.BOOL, wintypes.DWORD)
+ kernel32.SetConsoleCtrlHandler.argtypes = [HANDLER, wintypes.BOOL]
+ kernel32.SetConsoleCtrlHandler.restype = wintypes.BOOL
+
+ def handler(event):
+ if event == 0: # CTRL_C_EVENT
+ _thread.interrupt_main()
+ return True
+ return False
+
+ callback = HANDLER(handler)
+ if not kernel32.SetConsoleCtrlHandler(callback, True):
+ raise ctypes.WinError(ctypes.get_last_error())
+ try:
+ yield
+ finally:
+ if not kernel32.SetConsoleCtrlHandler(callback, False):
+ raise ctypes.WinError(ctypes.get_last_error())
+
+
+def _should_arm_external_stop() -> bool:
+ return (
+ threading.current_thread() is threading.main_thread()
+ and os.environ.get("PYTHON_JULIACALL_HANDLE_SIGNALS") == "yes"
+ )
+
+
+def _stop_channel(cleanup: ExitStack) -> tuple[int, int]:
+ """Open a nonblocking pair of descriptors carrying the signal wakeup byte.
+
+ Windows writes that byte with `send` and can only poll sockets, so a socket
+ pair is the one transport that works everywhere.
+ """
+ read_end, write_end = socket.socketpair()
+ cleanup.callback(read_end.close)
+ cleanup.callback(write_end.close)
+ read_end.setblocking(False)
+ write_end.setblocking(False)
+ return read_end.fileno(), write_end.fileno()
+
+
+@contextmanager
+def _external_stop_signal_context(model):
+ interrupted = False
+ external_stop = None
+
+ with ExitStack() as cleanup:
+ if _should_arm_external_stop():
+ stop_read_fd, stop_write_fd = _stop_channel(cleanup)
+ external_stop = SymbolicRegression.ExternalStop(stop_read_fd, signal.SIGINT)
+
+ # `signal.signal` below replaces the SIGINT handler Julia installed
+ # on POSIX, and sigaction is the only way to put it back.
+ if os.name == "posix":
+ libc = _libc_with_sigaction()
+ saved_sigaction = _SigactionStorage()
+ _checked_sigaction(libc, None, ctypes.byref(saved_sigaction))
+ cleanup.callback(
+ _checked_sigaction, libc, ctypes.byref(saved_sigaction), None
+ )
+
+ saved_python_handler = signal.getsignal(signal.SIGINT)
+
+ def record_interrupt(*_):
+ nonlocal interrupted
+ interrupted = True
+
+ signal.signal(signal.SIGINT, record_interrupt)
+ cleanup.callback(signal.signal, signal.SIGINT, saved_python_handler)
+ previous_wakeup_fd = signal.set_wakeup_fd(stop_write_fd)
+ cleanup.callback(signal.set_wakeup_fd, previous_wakeup_fd)
+
+ # On Windows, Julia takes Ctrl-C through a console handler instead.
+ # Registering ours last means it is removed first, while
+ # `record_interrupt` still absorbs a Ctrl-C arriving mid-teardown.
+ if os.name != "posix":
+ cleanup.enter_context(_console_ctrl_c_to_python())
+
+ yield external_stop
+
+ model.interrupted_ = interrupted
+ if interrupted:
+ warnings.warn(
+ "The search was interrupted. Returning partial results.",
+ RuntimeWarning,
+ stacklevel=3,
+ )
diff --git a/pysr/julia_extensions.py b/pysr/julia_extensions.py
index 6d16756e5..e9cd18200 100644
--- a/pysr/julia_extensions.py
+++ b/pysr/julia_extensions.py
@@ -16,6 +16,7 @@ def load_required_packages(
autodiff_backend: Literal["Zygote", "Mooncake", "Enzyme"] | None = None,
cluster_manager: str | None = None,
logger_spec: AbstractLoggerSpec | None = None,
+ use_tracing: bool = False,
):
if turbo:
load_package("LoopVectorization", "bdcacae8-1622-11e9-2a5c-532679323890")
@@ -27,10 +28,14 @@ def load_required_packages(
load_package("Mooncake", "da2b9cff-9c12-43a0-ae48-6db2b0edb7d6")
elif autodiff_backend == "Enzyme":
load_package("Enzyme", "7da242da-08ed-463a-9acd-ee780be4f1d9")
- if cluster_manager is not None:
+ if cluster_manager == "slurm":
+ load_package("SlurmClusterManager", "c82cd089-7bf7-41d7-976b-6b5d413cbe0a")
+ elif cluster_manager is not None:
load_package("ClusterManagers", "34f1f09b-3a8b-5176-ab39-66d58a4d544e")
if isinstance(logger_spec, TensorBoardLoggerSpec):
load_package("TensorBoardLogger", "899adc3e-224a-11e9-021f-63837185c80f")
+ if use_tracing:
+ load_package("JSON", "682c06a0-de6a-54ab-a142-c8b1cf79cde6")
def load_all_packages():
@@ -41,7 +46,9 @@ def load_all_packages():
autodiff_backend="Zygote",
cluster_manager="slurm",
logger_spec=TensorBoardLoggerSpec(log_dir="logs"),
+ use_tracing=True,
)
+ load_package("ClusterManagers", "34f1f09b-3a8b-5176-ab39-66d58a4d544e")
# TODO: Refactor this file so we can install all packages at once using `juliapkg`,
@@ -49,7 +56,9 @@ def load_all_packages():
def isinstalled(uuid_s: str):
- return jl.haskey(Pkg.dependencies(), jl.Base.UUID(uuid_s))
+ dependencies = Pkg.dependencies()
+ uuid = jl.Base.UUID(uuid_s)
+ return jl.haskey(dependencies, uuid) and dependencies[uuid].is_direct_dep
def load_package(package_name: str, uuid_s: str) -> None:
diff --git a/pysr/julia_helpers.py b/pysr/julia_helpers.py
index b72793105..f7411d1df 100644
--- a/pysr/julia_helpers.py
+++ b/pysr/julia_helpers.py
@@ -30,6 +30,19 @@ def _escape_filename(filename):
def _load_cluster_manager(cluster_manager: str):
+ if cluster_manager == "slurm":
+ jl.seval("using Distributed: addprocs")
+ jl.seval("using SlurmClusterManager: SlurmManager")
+ return jl.seval("""
+ (numprocs; kws...) -> begin
+ manager = SlurmManager()
+ manager.ntasks == numprocs || error(
+ "Requested $numprocs processes, but Slurm allocation has $(manager.ntasks) tasks. " *
+ "Set Slurm `--ntasks`/`--ntasks-per-node` and `procs` to the same value."
+ )
+ addprocs(manager; kws...)
+ end
+ """)
jl.seval(f"using ClusterManagers: addprocs_{cluster_manager}")
return jl.seval(f"addprocs_{cluster_manager}")
@@ -43,6 +56,10 @@ def jl_array(x, dtype=None):
return jl_convert(jl.Array[dtype], x)
+def jl_numpy_array(x):
+ return jl.copy(PythonCall.PyArray(x))
+
+
def jl_dict(x):
return jl_convert(jl.Dict, x)
@@ -52,7 +69,7 @@ def jl_named_tuple(d):
def jl_is_function(f) -> bool:
- return cast(bool, jl.seval("op -> op isa Function")(f))
+ return cast(bool, jl.isa(f, jl.Function))
def jl_serialize(obj: Any) -> NDArray[np.uint8]:
diff --git a/pysr/julia_import.py b/pysr/julia_import.py
index 0de6f499a..72c2610e4 100644
--- a/pysr/julia_import.py
+++ b/pysr/julia_import.py
@@ -6,6 +6,18 @@
from .julia_registry_helpers import try_with_registry_fallback
+autoload_extensions = os.environ.get("PYSR_AUTOLOAD_EXTENSIONS")
+if autoload_extensions is not None:
+ warnings.warn(
+ "PYSR_AUTOLOAD_EXTENSIONS is deprecated and will be removed in a future release. "
+ "Set PYTHON_JULIACALL_AUTOLOAD_IPYTHON_EXTENSION instead.",
+ FutureWarning,
+ )
+ # Only fill the gap; an explicitly set juliacall variable takes precedence.
+ os.environ.setdefault(
+ "PYTHON_JULIACALL_AUTOLOAD_IPYTHON_EXTENSION", autoload_extensions
+ )
+
# Check if JuliaCall is already loaded, and if so, warn the user
# about the relevant environment variables. If not loaded,
# set up sensible defaults.
@@ -35,16 +47,18 @@
("PYTHON_JULIACALL_HANDLE_SIGNALS", "yes"),
("PYTHON_JULIACALL_THREADS", "auto"),
("PYTHON_JULIACALL_OPTLEVEL", "3"),
+ # Don't hijack `%%julia` magics in notebooks unless asked;
+ # opt back in with PYTHON_JULIACALL_AUTOLOAD_IPYTHON_EXTENSION=yes
+ # or `%load_ext juliacall`.
+ ("PYTHON_JULIACALL_AUTOLOAD_IPYTHON_EXTENSION", "no"),
+ # Concurrent GC sweeping (`N,1`): the search is GC-bound at high
+ # thread counts, and this was neutral at 8 threads, 8% faster at 32,
+ # and 25% faster at 96.
+ ("JULIA_NUM_GC_THREADS", f"{max(1, (os.cpu_count() or 2) // 2)},1"),
):
os.environ[k] = os.environ.get(k, default)
-autoload_extensions = os.environ.get("PYSR_AUTOLOAD_EXTENSIONS")
-if autoload_extensions is not None:
- # Deprecated; so just pass to juliacall
- os.environ["PYTHON_JULIACALL_AUTOLOAD_IPYTHON_EXTENSION"] = autoload_extensions
-
-
def _import_juliacall():
import juliacall # type: ignore
diff --git a/pysr/juliapkg.json b/pysr/juliapkg.json
index 532a94a21..f786583e2 100644
--- a/pysr/juliapkg.json
+++ b/pysr/juliapkg.json
@@ -3,8 +3,10 @@
"packages": {
"SymbolicRegression": {
"uuid": "8254be44-1295-4e6a-a16d-46603ac705cb",
- "url": "https://github.com/astroautomata/SymbolicRegression.jl",
- "rev": "v2.0.0-alpha.9"
+ "version": "~2.5.0",
+ "preferences": {
+ "precompile_float64": false
+ }
},
"Serialization": {
"uuid": "9e88b42a-f829-5b0c-bbe9-9e923198166b",
diff --git a/pysr/mutations.py b/pysr/mutations.py
new file mode 100644
index 000000000..57cc5052c
--- /dev/null
+++ b/pysr/mutations.py
@@ -0,0 +1,187 @@
+"""Mutation configurations for :class:`PySRRegressor`."""
+
+from __future__ import annotations
+
+from abc import ABC, abstractmethod
+from collections.abc import Mapping, Sequence
+from dataclasses import dataclass
+
+from .julia_helpers import jl_array
+from .julia_import import AnyValue, SymbolicRegression, jl
+
+_LEGACY_MUTATION_PARAMETERS = (
+ "weight_add_node",
+ "weight_insert_node",
+ "weight_delete_node",
+ "weight_do_nothing",
+ "weight_mutate_constant",
+ "weight_mutate_operator",
+ "weight_mutate_feature",
+ "weight_swap_operands",
+ "weight_rotate_tree",
+ "weight_randomize",
+ "weight_simplify",
+ "weight_optimize",
+ "weight_backsolve",
+)
+
+
+def create_mutations(mutation_pairs: Sequence[tuple[AnyValue, float]]) -> AnyValue:
+ """Create a Julia vector of weighted mutations."""
+ pair_type = jl.Pair[SymbolicRegression.AbstractMutation, jl.Float64]
+ return jl_array(mutation_pairs, dtype=pair_type)
+
+
+def convert_mutations(
+ mutation_weights: Mapping[AbstractMutation, float],
+) -> AnyValue:
+ """Convert Python mutation weights to their Julia representation."""
+ return create_mutations(
+ [
+ jl.Pair(mutation.julia_mutation(), weight)
+ for mutation, weight in mutation_weights.items()
+ ]
+ )
+
+
+class AbstractMutation(ABC):
+ """Base class for mutation configurations."""
+
+ @abstractmethod
+ def julia_mutation(self) -> AnyValue:
+ """Create the corresponding SymbolicRegression.jl mutation."""
+ pass # pragma: no cover
+
+
+@dataclass(frozen=True)
+class _ParameterlessMutation(AbstractMutation):
+ def julia_mutation(self) -> AnyValue:
+ return getattr(SymbolicRegression, type(self).__name__)()
+
+
+@dataclass(frozen=True)
+class ConstantMutation(AbstractMutation):
+ """Perturb a constant.
+
+ SymbolicRegression.jl default weight: ``0.0346``.
+
+ Defaults match SymbolicRegression.jl.
+ """
+
+ perturbation_factor: float = 0.086
+ probability_negate: float = 0.01
+
+ def julia_mutation(self) -> AnyValue:
+ return SymbolicRegression.ConstantMutation(
+ perturbation_factor=self.perturbation_factor,
+ probability_negate=self.probability_negate,
+ )
+
+
+@dataclass(frozen=True)
+class OperatorMutation(_ParameterlessMutation):
+ """Replace an operator with another operator of the same arity.
+
+ SymbolicRegression.jl default weight: ``0.293``.
+ """
+
+
+@dataclass(frozen=True)
+class FeatureMutation(_ParameterlessMutation):
+ """Change the feature referenced by a variable node.
+
+ SymbolicRegression.jl default weight: ``0.1``.
+ """
+
+
+@dataclass(frozen=True)
+class SwapOperandsMutation(_ParameterlessMutation):
+ """Swap the operands of a binary operator.
+
+ SymbolicRegression.jl default weight: ``0.198``.
+ """
+
+
+@dataclass(frozen=True)
+class AddNodeMutation(_ParameterlessMutation):
+ """Append a node to the expression.
+
+ SymbolicRegression.jl default weight: ``2.47``.
+ """
+
+
+@dataclass(frozen=True)
+class InsertNodeMutation(_ParameterlessMutation):
+ """Insert a node above an existing node.
+
+ SymbolicRegression.jl default weight: ``0.0112``.
+ """
+
+
+@dataclass(frozen=True)
+class DeleteNodeMutation(_ParameterlessMutation):
+ """Delete a node from the expression.
+
+ SymbolicRegression.jl default weight: ``0.870``.
+ """
+
+
+@dataclass(frozen=True)
+class RotateTreeMutation(_ParameterlessMutation):
+ """Rotate a subtree.
+
+ SymbolicRegression.jl default weight: ``4.26``.
+ """
+
+
+@dataclass(frozen=True)
+class BacksolveMutation(AbstractMutation):
+ """Fit a replacement expression by backsolving through the expression.
+
+ SymbolicRegression.jl default weight: ``0.0``.
+ """
+
+ max_library_size: int = 500
+ max_terms: int = 8
+ min_improvement: float = 1e-3
+ node_attempts: int = 8
+
+ def julia_mutation(self) -> AnyValue:
+ return SymbolicRegression.BacksolveMutation(
+ max_library_size=self.max_library_size,
+ max_terms=self.max_terms,
+ min_improvement=self.min_improvement,
+ node_attempts=self.node_attempts,
+ )
+
+
+@dataclass(frozen=True)
+class SimplifyMutation(_ParameterlessMutation):
+ """Simplify constant parts of the expression.
+
+ SymbolicRegression.jl default weight: ``0.00209``.
+ """
+
+
+@dataclass(frozen=True)
+class RandomizeMutation(_ParameterlessMutation):
+ """Replace the expression with a random expression.
+
+ SymbolicRegression.jl default weight: ``0.000502``.
+ """
+
+
+@dataclass(frozen=True)
+class OptimizeMutation(_ParameterlessMutation):
+ """Optimize constants as a mutation.
+
+ SymbolicRegression.jl default weight: ``0.0``.
+ """
+
+
+@dataclass(frozen=True)
+class DoNothingMutation(_ParameterlessMutation):
+ """Leave the expression unchanged.
+
+ SymbolicRegression.jl default weight: ``0.273``.
+ """
diff --git a/pysr/param_groupings.yml b/pysr/param_groupings.yml
index f4b769c27..70cd1b93c 100644
--- a/pysr/param_groupings.yml
+++ b/pysr/param_groupings.yml
@@ -4,6 +4,7 @@
- unary_operators
- operators
- expression_spec
+ - type_spec
- maxsize
- maxdepth
- Setting the Search Size:
@@ -33,6 +34,10 @@
- adaptive_parsimony_scaling
- should_simplify
- Mutations:
+ - default_mutations
+ - mutations
+ - default_plugins
+ - plugins
- weight_add_node
- weight_insert_node
- weight_delete_node
@@ -45,6 +50,7 @@
- weight_randomize
- weight_simplify
- weight_optimize
+ - weight_backsolve
- crossover_probability
- annealing
- alpha
@@ -101,6 +107,8 @@
- progress
- logger_spec
- input_stream
+ - use_tracing
+ - tracing_file
- Environment:
- temp_equation_file
- tempdir
diff --git a/pysr/plugins.py b/pysr/plugins.py
new file mode 100644
index 000000000..ade737c97
--- /dev/null
+++ b/pysr/plugins.py
@@ -0,0 +1,89 @@
+"""Plugin configurations for :class:`PySRRegressor`."""
+
+from __future__ import annotations
+
+from abc import ABC, abstractmethod
+from dataclasses import dataclass
+from typing import Literal
+
+from .julia_import import AnyValue, SymbolicRegression, jl
+
+
+class AbstractPlugin(ABC):
+ """Base class for plugin configurations."""
+
+ @abstractmethod
+ def julia_plugin(self) -> AnyValue:
+ """Create the corresponding SymbolicRegression.jl plugin."""
+ pass # pragma: no cover
+
+
+@dataclass(frozen=True)
+class SimulatedAnnealingPlugin(AbstractPlugin):
+ """Apply simulated annealing during mutation acceptance.
+
+ Enabled by default through ``annealing=True``.
+ """
+
+ alpha: float = 0.1
+
+ def julia_plugin(self) -> AnyValue:
+ return SymbolicRegression.SimulatedAnnealingPlugin(alpha=self.alpha)
+
+
+@dataclass(frozen=True)
+class AdaptiveParsimonyPlugin(AbstractPlugin):
+ """Apply frequency-based parsimony during selection and mutation.
+
+ Enabled by default in PySR through ``use_frequency=True`` and
+ ``use_frequency_in_tournament=True``.
+ """
+
+ tournament: bool = True
+ mutation_acceptance: bool = True
+
+ def julia_plugin(self) -> AnyValue:
+ return SymbolicRegression.AdaptiveParsimonyPlugin(
+ tournament=self.tournament,
+ mutation_acceptance=self.mutation_acceptance,
+ )
+
+
+@dataclass(frozen=True)
+class AdaptiveMutationWeightsPlugin(AbstractPlugin):
+ """Adapt mutation weights from their observed success rates.
+
+ Enabled by default in PySR. Pass an explicit ``default_plugins`` that
+ omits this plugin to disable it.
+ Learned runtime state is reinitialized for each call to ``fit``.
+ """
+
+ smoothing: float = 0.02
+ floor: float = 0.05
+ reward: Literal["cost", "loss"] = "cost"
+
+ def julia_plugin(self) -> AnyValue:
+ return SymbolicRegression.AdaptiveMutationWeightsPlugin(
+ smoothing=self.smoothing,
+ floor=self.floor,
+ reward=jl.Symbol(self.reward),
+ )
+
+
+@dataclass(frozen=True)
+class MutationBurstPlugin(AbstractPlugin):
+ """Retry rejected mutations and optionally chain accepted mutations.
+
+ Disabled by default in PySR.
+ """
+
+ retry_attempts: int = 4
+ compound_probability: float = 0.25
+ compound_max_steps: int = 2
+
+ def julia_plugin(self) -> AnyValue:
+ return SymbolicRegression.MutationBurstPlugin(
+ retry_attempts=self.retry_attempts,
+ compound_probability=self.compound_probability,
+ compound_max_steps=self.compound_max_steps,
+ )
diff --git a/pysr/sr.py b/pysr/sr.py
index f09b8bc50..fa4b46d6f 100644
--- a/pysr/sr.py
+++ b/pysr/sr.py
@@ -10,8 +10,9 @@
import sys
import tempfile
import warnings
-from collections.abc import Callable
+from collections.abc import Callable, Mapping, Sequence
from dataclasses import dataclass, fields
+from functools import wraps
from io import StringIO
from multiprocessing import cpu_count
from pathlib import Path
@@ -38,10 +39,10 @@
from .expression_specs import (
AbstractExpressionSpec,
ExpressionSpec,
- ParametricExpressionSpec,
- parametric_expression_deprecation_warning,
+ TemplateExpressionSpec,
)
from .feature_selection import run_feature_selection
+from .interrupt import _external_stop_signal_context
from .julia_extensions import load_required_packages
from .julia_helpers import (
_escape_filename,
@@ -50,10 +51,33 @@
jl_deserialize,
jl_is_function,
jl_named_tuple,
+ jl_numpy_array,
jl_serialize,
)
from .julia_import import AnyValue, SymbolicRegression, VectorValue, jl
from .logger_specs import AbstractLoggerSpec
+from .mutations import (
+ _LEGACY_MUTATION_PARAMETERS,
+ AbstractMutation,
+ convert_mutations,
+)
+from .plugins import AbstractPlugin
+from .type_specs import (
+ TypeSpec,
+ _TypeSpecRuntime,
+ _TypeSpecRuntimeDefinition,
+ compile_type_spec_runtime_for_model,
+ create_type_spec_addprocs_function,
+ create_type_spec_exports,
+ load_type_spec_runtime,
+ object_array_1d,
+ prepare_type_spec_fit_data,
+ prepare_type_spec_prediction_data,
+ type_spec_to_julia_array,
+ validate_type_spec_model_configuration,
+ validate_type_spec_options,
+ validate_type_spec_runtime,
+)
from .utils import (
ArrayLike,
PathLike,
@@ -76,8 +100,11 @@
from typing_extensions import List
+_CHECKPOINT_SCHEMA_VERSION = 3
+
ALREADY_RAN = False
+
pysr_logger = logging.getLogger(__name__)
@@ -103,7 +130,7 @@ def _process_constraints(
"One typical constraint is to use `constraints={..., '^': (-1, 1)}`, which "
"will allow arbitrary-complexity base (-1) but only powers such as "
"a constant or variable (1). "
- "For more tips, please see https://ai.damtp.cam.ac.uk/pysr/tuning/"
+ "For more tips, please see https://pysr.ai/tuning"
)
# Create default constraint tuple with -1 for each argument
constraints[op] = tuple([-1] * arity)
@@ -141,7 +168,7 @@ def _process_constraints(
def _maybe_create_inline_operators(
operators: dict[int, list[str]],
extra_sympy_mappings: dict[str, Callable] | None,
- expression_spec: AbstractExpressionSpec,
+ supports_sympy: bool,
) -> dict[int, list[str]]:
operators = {arity: op_list.copy() for arity, op_list in operators.items()}
@@ -150,24 +177,22 @@ def _maybe_create_inline_operators(
is_user_defined_operator = "(" in op
if is_user_defined_operator:
- jl.seval(op)
- # Cut off from the first non-alphanumeric char:
- first_non_char = [j for j, char in enumerate(op) if char == "("][0]
- function_name = op[:first_non_char]
- # Assert that function_name only contains
- # alphabetical characters, numbers,
- # and underscores:
+ function = jl.seval(op)
+ if not jl_is_function(function):
+ raise ValueError(
+ "Custom operator definitions must evaluate to a Julia function."
+ )
+ function_name = str(jl.Base.nameof(function))
if not re.match(r"^[a-zA-Z0-9_]+$", function_name):
raise ValueError(
- f"Invalid function name {function_name}. "
- "Only alphanumeric characters, numbers, "
- "and underscores are allowed."
+ "Custom operators must define a named Julia function whose name "
+ "contains only alphanumeric characters, numbers, and underscores."
)
missing_sympy_mapping = (
extra_sympy_mappings is None
or function_name not in extra_sympy_mappings
)
- if missing_sympy_mapping and expression_spec.supports_sympy:
+ if missing_sympy_mapping and supports_sympy:
raise ValueError(
f"Custom function {function_name} is not defined in `extra_sympy_mappings`. "
"You can define it with, "
@@ -179,6 +204,41 @@ def _maybe_create_inline_operators(
return operators
+def _create_julia_operators_and_loss_functions(
+ operators: dict[int, list[str]],
+ extra_sympy_mappings: dict[str, Callable] | None,
+ supports_sympy: bool,
+ elementwise_loss: str | None,
+ loss_function: str | None,
+ loss_function_expression: str | None,
+) -> tuple[dict[int, list[str]], AnyValue | None, AnyValue | None, AnyValue | None]:
+ operators = _maybe_create_inline_operators(
+ operators=operators,
+ extra_sympy_mappings=extra_sympy_mappings,
+ supports_sympy=supports_sympy,
+ )
+
+ def eval_source(source: str) -> Any:
+ return jl.seval(source)
+
+ def eval_objective(source: str | None, knob: str) -> AnyValue | None:
+ if source is None:
+ return None
+ loss = eval_source(str(source))
+ if not jl_is_function(loss):
+ raise ValueError(f"`{knob}` must evaluate to a callable Julia function.")
+ return loss
+
+ custom_loss = (
+ None if elementwise_loss is None else eval_source(str(elementwise_loss))
+ )
+ custom_full_objective = eval_objective(loss_function, "loss_function")
+ custom_loss_expression = eval_objective(
+ loss_function_expression, "loss_function_expression"
+ )
+ return operators, custom_loss, custom_full_objective, custom_loss_expression
+
+
def _check_assertions(
X,
use_custom_variable_names,
@@ -188,6 +248,7 @@ def _check_assertions(
y,
X_units,
y_units,
+ supports_sympy,
):
# Check for potential errors before they happen
assert len(X.shape) == 2
@@ -197,7 +258,8 @@ def _check_assertions(
assert weights.shape == y.shape
assert X.shape[0] == weights.shape[0]
if use_custom_variable_names:
- assert len(variable_names) == X.shape[1]
+ if len(variable_names) != X.shape[1]:
+ raise ValueError("`variable_names` must contain one name per feature.")
# Check none of the variable names are function names:
for var_name in variable_names:
# Check if alphanumeric only:
@@ -207,7 +269,8 @@ def _check_assertions(
"Only alphanumeric characters, numbers, "
"and underscores are allowed."
)
- assert_valid_sympy_symbol(var_name)
+ if supports_sympy:
+ assert_valid_sympy_symbol(var_name)
if (
isinstance(complexity_of_variables, list)
and len(complexity_of_variables) != X.shape[1]
@@ -343,6 +406,41 @@ def _validate_export_mappings(extra_jax_mappings, extra_torch_mappings):
# Class validation constants
VALID_OPTIMIZER_ALGORITHMS = ["BFGS", "NelderMead"]
+_WARM_START_MUTABLE_STATE = (
+ "equations_",
+ "equation_file_contents_",
+ "julia_state_stream_",
+ "julia_options_stream_",
+ "selection_mask_",
+ "feature_names_in_",
+ "display_feature_names_in_",
+ "complexity_of_variables_",
+ "X_units_",
+ "y_units_",
+)
+
+
+def _rollback_failed_warm_start(
+ fit_method: Callable[..., Any],
+) -> Callable[..., Any]:
+ @wraps(fit_method)
+ def wrapped(self, *args, **kwargs):
+ previous_state = None
+ if self.warm_start and getattr(self, "julia_state_stream_", None) is not None:
+ previous_state = self.__dict__.copy()
+ for name in _WARM_START_MUTABLE_STATE:
+ if name in previous_state:
+ previous_state[name] = copy.deepcopy(previous_state[name])
+ try:
+ return fit_method(self, *args, **kwargs)
+ except BaseException:
+ if previous_state is not None:
+ self.__dict__.clear()
+ self.__dict__.update(previous_state)
+ raise
+
+ return wrapped
+
@dataclass
class _DynamicallySetParams:
@@ -369,7 +467,7 @@ class PySRRegressor(MultiOutputMixin, RegressorMixin, BaseEstimator):
Most default parameters have been tuned over several example equations,
but you should adjust `niterations`, `binary_operators`, `unary_operators`
to your requirements. You can view more detailed explanations of the options
- on the [options page](https://ai.damtp.cam.ac.uk/pysr/options) of the
+ on the [options page](https://pysr.ai/options) of the
documentation.
Parameters
@@ -389,7 +487,7 @@ class PySRRegressor(MultiOutputMixin, RegressorMixin, BaseEstimator):
most accurate model.
binary_operators : list[str]
List of strings for binary operators used in the search.
- See the [operators page](https://ai.damtp.cam.ac.uk/pysr/operators/)
+ See the [operators page](https://pysr.ai/operators)
for more details.
Default is `["+", "-", "*", "/"]`.
unary_operators : list[str]
@@ -408,6 +506,10 @@ class PySRRegressor(MultiOutputMixin, RegressorMixin, BaseEstimator):
`TemplateExpressionSpec(...)` which allows you to specify
a custom template for the expressions.
Default is `ExpressionSpec()`.
+ type_spec : TypeSpec
+ Declarative custom value type. PySR builds an isolated Julia wrapper,
+ interface methods, operators, and loss from this specification.
+ Default is `None`, which uses the numeric path.
niterations : int
Number of iterations of the algorithm to run. The best
equations are printed and migrate between populations at the
@@ -567,7 +669,7 @@ class PySRRegressor(MultiOutputMixin, RegressorMixin, BaseEstimator):
(requires `annealing` to be `True`).
Default is `3.17`.
annealing : bool
- Whether to use annealing. Default is `False`.
+ Whether to use annealing. Default is `True`.
early_stop_condition : float | str
Stop the search early if this loss is reached. You may also
pass a string containing a Julia function which
@@ -588,50 +690,61 @@ class PySRRegressor(MultiOutputMixin, RegressorMixin, BaseEstimator):
fraction_replaced_guesses : float
How much of the population to replace with migrating equations from
guesses. Default is `0.001`.
- weight_add_node : float
+ weight_add_node : float | None
Relative likelihood for mutation to add a node.
- Default is `2.47`.
- weight_insert_node : float
+ Default is `None` (mapping to `2.47`).
+ weight_insert_node : float | None
Relative likelihood for mutation to insert a node.
- Default is `0.0112`.
- weight_delete_node : float
+ Default is `None` (mapping to `0.0112`).
+ weight_delete_node : float | None
Relative likelihood for mutation to delete a node.
- Default is `0.870`.
- weight_do_nothing : float
+ Default is `None` (mapping to `0.870`).
+ weight_do_nothing : float | None
Relative likelihood for mutation to leave the individual.
- Default is `0.273`.
- weight_mutate_constant : float
+ Default is `None` (mapping to `0.273`).
+ weight_mutate_constant : float | None
Relative likelihood for mutation to change the constant slightly
in a random direction.
- Default is `0.0346`.
- weight_mutate_operator : float
+ Default is `None` (mapping to `0.0346`).
+ weight_mutate_operator : float | None
Relative likelihood for mutation to swap an operator.
- Default is `0.293`.
- weight_mutate_feature : float
+ Default is `None` (mapping to `0.293`).
+ weight_mutate_feature : float | None
Relative likelihood for mutation to change which feature a variable node references.
- Default is `0.1`.
- weight_swap_operands : float
+ Default is `None` (mapping to `0.1`).
+ weight_swap_operands : float | None
Relative likehood for swapping operands in binary operators.
- Default is `0.198`.
- weight_rotate_tree : float
+ Default is `None` (mapping to `0.198`).
+ weight_rotate_tree : float | None
How often to perform a tree rotation at a random node.
- Default is `4.26`.
- weight_randomize : float
+ Default is `None` (mapping to `4.26`).
+ weight_randomize : float | None
Relative likelihood for mutation to completely delete and then
randomly generate the equation
- Default is `0.000502`.
- weight_simplify : float
+ Default is `None` (mapping to `0.000502`).
+ weight_simplify : float | None
Relative likelihood for mutation to simplify constant parts by evaluation
- Default is `0.00209`.
- weight_optimize: float
+ Default is `None` (mapping to `0.00209`).
+ weight_optimize: float | None
Constant optimization can also be performed as a mutation, in addition to
the normal strategy controlled by `optimize_probability` which happens
every iteration. Using it as a mutation is useful if you want to use
a large `ncycles_periteration`, and may not optimize very often.
- Default is `0.0`.
+ Default is `None` (mapping to `0.0`).
+ weight_backsolve : float | None
+ Relative likelihood for backsolve mutation. To configure its parameters,
+ pass `BacksolveMutation(...)` through `mutations` instead.
+ Default is `None` (mapping to `0.0`).
+ mutations : Mapping[AbstractMutation, float] | None
+ Mutation configurations and their weights. Entries override or extend
+ `default_mutations` by mutation type. Default is `None`.
+ default_mutations : Mapping[AbstractMutation, float] | None
+ Default mutation configurations and their weights. When provided, these
+ replace the SymbolicRegression.jl defaults and plugin-contributed mutations.
+ Default is `None`.
crossover_probability : float
Absolute probability of crossover-type genetic operation, instead of a mutation.
- Default is `0.0259`.
+ Default is `0.2`.
skip_mutation_failures : bool
Whether to skip mutation and crossover failures, rather than
simply re-sampling the current member.
@@ -690,11 +803,10 @@ class PySRRegressor(MultiOutputMixin, RegressorMixin, BaseEstimator):
Number of processes to use for parallelism. If `None`, defaults to `cpu_count()`.
Default is `None`.
cluster_manager : str
- For distributed computing, this sets the job queue system. Set
- to one of "slurm", "pbs", "lsf", "sge", "qrsh", "scyld", or
- "htc". If set to one of these, PySR will run in distributed
- mode, and use `procs` to figure out how many processes to launch.
- Default is `None`.
+ For distributed computing, this sets the job queue system. Set to
+ "slurm" to use an existing Slurm allocation, with `procs` equal to
+ its task count. Other supported values are "pbs", "lsf", "sge",
+ "qrsh", "scyld", and "htc". Default is `None`.
heap_size_hint_in_bytes : int
For multiprocessing, this sets the `--heap-size-hint` parameter
for new Julia processes. This can be configured when using
@@ -706,13 +818,17 @@ class PySRRegressor(MultiOutputMixin, RegressorMixin, BaseEstimator):
If a worker does not respond within this time, it will be restarted.
Default is `None`.
worker_imports : list[str] | None
- List of module names as strings to import in worker processes.
- For example, `["MyPackage", "OtherPackage"]` will run `using MyPackage, OtherPackage`
- in each worker process. Default is `None`.
+ Module names to import on multiprocessing workers. For example,
+ ``["MyPackage", "OtherPackage"]`` runs
+ ``using MyPackage, OtherPackage`` on every worker. TypeSpec preambles,
+ operators, objectives, and expression specifications may depend on
+ packages listed here; arbitrary bindings from ``Main`` are unavailable.
+ Default is ``None``.
batching : bool | "auto"
Whether to compare population members on small batches during
evolution. Still uses full dataset for comparing against hall
- of fame. "auto" enables batching for N>1000. Default is `"auto"`.
+ of fame. `"auto"` lets SymbolicRegression.jl choose based on the
+ dataset. Default is `"auto"`.
batch_size : int | None
The batch size to use if batching. If None, uses the
full dataset when N<=1000, 128 for N<5000, 256 for N<50000,
@@ -756,12 +872,17 @@ class PySRRegressor(MultiOutputMixin, RegressorMixin, BaseEstimator):
warm_start : bool
Tells fit to continue from where the last call to fit finished.
If false, each call to fit will be fresh, overwriting previous results.
+ Plugin runtime state is reinitialized for each call to `fit`.
Default is `False`.
- guesses : list[str] | list[list[str]] | list[dict[str, str]] | list[list[dict[str, str]]] | None
+ guesses : list[str] | list[list[str]] | list[dict[str, str | ArrayLike]] | list[list[dict[str, str | ArrayLike]]] | None
Initial guesses for expressions to seed the search. Examples:
`["x0 + x1", "x0^2"]`, `[["x0"], ["x1"]]` (multi-output),
- `[{"f": "#1 + #2"}]` (TemplateExpressionSpec where `#1`, `#2` are
- placeholders for the 1st, 2nd arguments of expression `f`).
+ or `[{"f": "#1 + #2", "p": [5.0, 10.0]}]` for a
+ `TemplateExpressionSpec` with component `f` and parameter vector `p`.
+ Here `#1` and `#2` are the first and second arguments of `f`.
+ Parameter values use the data precision or the model's `TypeSpec`, and
+ supplied vectors must match the lengths declared by the template.
+ Parameter names may be omitted to retain their normal initialization.
Default is `None`.
verbosity : int
What verbosity level to use. 0 means minimal print statements.
@@ -779,12 +900,30 @@ class PySRRegressor(MultiOutputMixin, RegressorMixin, BaseEstimator):
Logger specification for the Julia backend. See, for example,
`TensorBoardLoggerSpec`.
Default is `None`.
- input_stream : str
- The stream to read user input from. By default, this is `"stdin"`.
- If you encounter issues with reading from `stdin`, like a hang,
- you can simply pass `"devnull"` to this argument. You can also
- reference an arbitrary Julia object in the `Main` namespace.
- Default is `"stdin"`.
+ plugins : Sequence[AbstractPlugin] | None
+ Plugin configurations. Entries override or extend `default_plugins` by
+ plugin type. Default is `None`.
+ default_plugins : Sequence[AbstractPlugin] | None
+ Default plugin configurations. Default is `None`.
+ input_stream : str | None
+ The stream to read user input from, used for the `'q'` + ``
+ command that stops a search early. `None` reads from `"stdin"` when
+ attached to an interactive terminal and otherwise disables stdin
+ watching via `"devnull"` (for example in Jupyter, where typed input
+ never reaches the search). You can also pass `"stdin"` or `"devnull"`
+ explicitly, or reference an arbitrary Julia object in the `Main`
+ namespace. Default is `None`.
+ use_tracing : bool
+ Whether to write a JSONL trace of the search to `tracing_file`.
+ One record per line holds the live members of one population at one
+ iteration, along with every mutation, crossover, tuning and death
+ event. Uses JSON.jl, which is installed into the Julia environment
+ on first use.
+ Default is `False`.
+ tracing_file : str | Path
+ Where to write the trace when `use_tracing` is set. Any existing
+ file is overwritten when the search starts.
+ Default is `"pysr_trace.jsonl"`.
run_id : str
A unique identifier for the run. Will be generated using the
current date and time if not provided.
@@ -944,6 +1083,7 @@ class PySRRegressor(MultiOutputMixin, RegressorMixin, BaseEstimator):
logger_: AnyValue | None
equation_file_contents_: list[pd.DataFrame] | None
show_pickle_warnings_: bool
+ _type_spec_runtime_definition_: _TypeSpecRuntimeDefinition
def __init__(
self,
@@ -953,6 +1093,7 @@ def __init__(
unary_operators: list[str] | None = None,
operators: dict[int, list[str]] | None = None,
expression_spec: AbstractExpressionSpec | None = None,
+ type_spec: TypeSpec | None = None,
niterations: int = 100,
populations: int = 31,
population_size: int = 27,
@@ -978,25 +1119,28 @@ def __init__(
use_frequency_in_tournament: bool = True,
adaptive_parsimony_scaling: float = 1040.0,
alpha: float = 3.17,
- annealing: bool = False,
+ annealing: bool = True,
early_stop_condition: float | str | None = None,
ncycles_per_iteration: int = 380,
fraction_replaced: float = 0.00036,
fraction_replaced_hof: float = 0.0614,
fraction_replaced_guesses: float = 0.001,
- weight_add_node: float = 2.47,
- weight_insert_node: float = 0.0112,
- weight_delete_node: float = 0.870,
- weight_do_nothing: float = 0.273,
- weight_mutate_constant: float = 0.0346,
- weight_mutate_operator: float = 0.293,
- weight_mutate_feature: float = 0.1,
- weight_swap_operands: float = 0.198,
- weight_rotate_tree: float = 4.26,
- weight_randomize: float = 0.000502,
- weight_simplify: float = 0.00209,
- weight_optimize: float = 0.0,
- crossover_probability: float = 0.0259,
+ weight_add_node: float | None = None,
+ weight_insert_node: float | None = None,
+ weight_delete_node: float | None = None,
+ weight_do_nothing: float | None = None,
+ weight_mutate_constant: float | None = None,
+ weight_mutate_operator: float | None = None,
+ weight_mutate_feature: float | None = None,
+ weight_swap_operands: float | None = None,
+ weight_rotate_tree: float | None = None,
+ weight_randomize: float | None = None,
+ weight_simplify: float | None = None,
+ weight_optimize: float | None = None,
+ weight_backsolve: float | None = None,
+ mutations: Mapping[AbstractMutation, float] | None = None,
+ default_mutations: Mapping[AbstractMutation, float] | None = None,
+ crossover_probability: float = 0.2,
skip_mutation_failures: bool = True,
migration: bool = True,
hof_migration: bool = True,
@@ -1035,8 +1179,10 @@ def __init__(
guesses: (
list[str]
| list[list[str]]
- | list[dict[str, str]]
- | list[list[dict[str, str]]]
+ | list[dict[str, str | ArrayLike]]
+ | list[list[dict[str, str | ArrayLike]]]
+ | list[AnyValue]
+ | list[list[AnyValue]]
| None
) = None,
verbosity: int = 1,
@@ -1044,7 +1190,11 @@ def __init__(
print_precision: int = 5,
progress: bool = True,
logger_spec: AbstractLoggerSpec | None = None,
- input_stream: str = "stdin",
+ plugins: Sequence[AbstractPlugin] | None = None,
+ default_plugins: Sequence[AbstractPlugin] | None = None,
+ input_stream: str | None = None,
+ use_tracing: bool = False,
+ tracing_file: str | Path = "pysr_trace.jsonl",
run_id: str | None = None,
output_directory: str | None = None,
temp_equation_file: bool = False,
@@ -1067,6 +1217,7 @@ def __init__(
self.unary_operators = unary_operators
self.operators = operators
self.expression_spec = expression_spec
+ self.type_spec = type_spec
self.niterations = niterations
self.populations = populations
self.population_size = population_size
@@ -1113,6 +1264,9 @@ def __init__(
self.weight_randomize = weight_randomize
self.weight_simplify = weight_simplify
self.weight_optimize = weight_optimize
+ self.weight_backsolve = weight_backsolve
+ self.mutations = mutations
+ self.default_mutations = default_mutations
self.crossover_probability = crossover_probability
self.skip_mutation_failures = skip_mutation_failures
# -- Migration parameters
@@ -1159,7 +1313,11 @@ def __init__(
self.print_precision = print_precision
self.progress = progress
self.logger_spec = logger_spec
+ self.plugins = plugins
+ self.default_plugins = default_plugins
self.input_stream = input_stream
+ self.use_tracing = use_tracing
+ self.tracing_file = tracing_file
# - Project management
self.run_id = run_id
self.output_directory = output_directory
@@ -1202,7 +1360,7 @@ def __init__(
elif k == "julia_project":
warnings.warn(
"The `julia_project` parameter has been deprecated. To use a custom "
- "julia project, please see `https://ai.damtp.cam.ac.uk/pysr/backend`.",
+ "julia project, please see `https://pysr.ai/backend`.",
FutureWarning,
)
elif k == "julia_kwargs":
@@ -1301,11 +1459,19 @@ def from_file(
if "equations_" not in model.__dict__ or model.equations_ is None:
model.refresh()
+ else:
+ model._restore_julia_backed_columns()
- if model.expression_spec is not None:
+ if (
+ not isinstance(
+ model.expression_spec_, (ExpressionSpec, TemplateExpressionSpec)
+ )
+ and not model._has_fitted_type_spec()
+ ):
warnings.warn(
- "Loading model from checkpoint file with a non-default expression spec "
- "is not fully supported as it relies on dynamic objects. This may result in unexpected behavior.",
+ "Loading a checkpoint with a custom expression spec is not fully "
+ "supported, as it relies on dynamic Julia objects. This may result "
+ "in unexpected behavior.",
)
return model
@@ -1314,6 +1480,11 @@ def from_file(
f"Checkpoint file {pkl_filename} does not exist. "
"Attempting to recreate model from scratch..."
)
+ if pysr_kwargs.get("type_spec") is not None:
+ raise ValueError(
+ "TypeSpec models require the original `checkpoint.pkl`; "
+ "CSV-only reconstruction is not supported."
+ )
csv_filename = Path(run_directory) / "hall_of_fame.csv"
csv_filename_bak = Path(run_directory) / "hall_of_fame.csv.bak"
if not csv_filename.exists() and not csv_filename_bak.exists():
@@ -1436,27 +1607,33 @@ def __getstate__(self) -> dict[str, Any]:
warnings.warn(warn_msg)
else:
pysr_logger.debug(warn_msg)
- state_keys_to_clear = state_keys_containing_lambdas
- state_keys_to_clear.append("logger_")
+ state_keys_to_clear = [*state_keys_containing_lambdas, "logger_"]
pickled_state = {
key: (None if key in state_keys_to_clear else value)
for key, value in state.items()
}
+ pickled_state["_checkpoint_schema_version"] = _CHECKPOINT_SCHEMA_VERSION
if ("equations_" in pickled_state) and (
pickled_state["equations_"] is not None
):
pickled_state["output_torch_format"] = False
pickled_state["output_jax_format"] = False
+ unpicklable_columns = ["jax_format", "torch_format"]
+ if self._has_julia_backed_equations():
+ # Live Julia objects cannot be unpickled in a fresh process
+ # before their Julia definitions exist; these columns are
+ # rebuilt from `julia_state_` via `refresh()`.
+ unpicklable_columns += ["julia_expression", "lambda_format"]
if self.nout_ == 1:
pickled_columns = ~pickled_state["equations_"].columns.isin(
- ["jax_format", "torch_format"]
+ unpicklable_columns
)
pickled_state["equations_"] = (
pickled_state["equations_"].loc[:, pickled_columns].copy()
)
else:
pickled_columns = [
- ~dataframe.columns.isin(["jax_format", "torch_format"])
+ ~dataframe.columns.isin(unpicklable_columns)
for dataframe in pickled_state["equations_"]
]
pickled_state["equations_"] = [
@@ -1467,19 +1644,41 @@ def __getstate__(self) -> dict[str, Any]:
]
return pickled_state
+ def __setstate__(self, state: dict[str, Any]) -> None:
+ schema_version = state.pop("_checkpoint_schema_version", None)
+ state.setdefault("type_spec", None)
+ if schema_version != _CHECKPOINT_SCHEMA_VERSION:
+ raise ValueError(
+ "Unsupported PySR checkpoint schema: "
+ f"expected {_CHECKPOINT_SCHEMA_VERSION}, found {schema_version!r}."
+ )
+ self.__dict__.update(state)
+
def _checkpoint(self):
"""Save the model's current state to a checkpoint file.
This should only be used internally by PySRRegressor.
"""
- # Save model state:
+ checkpoint_path = self.get_pkl_filename()
+ previous_show_pickle_warnings = getattr(self, "show_pickle_warnings_", True)
+ # Same directory as the destination, so `os.replace` stays atomic and a
+ # failed pickle cannot truncate an existing checkpoint:
+ temporary_path = checkpoint_path.with_name(checkpoint_path.name + ".tmp")
self.show_pickle_warnings_ = False
- with open(self.get_pkl_filename(), "wb") as f:
+ try:
+ with open(temporary_path, "wb") as checkpoint_file:
+ try:
+ pkl.dump(self, checkpoint_file)
+ except Exception as e:
+ pysr_logger.debug(f"Error checkpointing model: {e}")
+ return
+ os.replace(temporary_path, checkpoint_path)
+ finally:
+ self.show_pickle_warnings_ = previous_show_pickle_warnings
try:
- pkl.dump(self, f)
+ temporary_path.unlink(missing_ok=True)
except Exception as e:
- pysr_logger.debug(f"Error checkpointing model: {e}")
- self.show_pickle_warnings_ = True
+ pysr_logger.debug(f"Error cleaning up temporary checkpoint file: {e}")
def get_pkl_filename(self) -> Path:
path = Path(self.output_directory_) / self.run_id_ / "checkpoint.pkl"
@@ -1498,14 +1697,22 @@ def equations(self): # pragma: no cover
@property
def julia_options_(self):
"""The deserialized julia options."""
- return jl_deserialize(self.julia_options_stream_)
+ stream = getattr(self, "julia_options_stream_", None)
+ if stream is None:
+ return None
+ self._define_julia_expression_types()
+ return jl_deserialize(stream)
@property
def julia_state_(self):
"""The deserialized state."""
+ stream = getattr(self, "julia_state_stream_", None)
+ if stream is None:
+ return None
+ self._define_julia_expression_types()
return cast(
Union[Tuple[VectorValue, AnyValue], None],
- jl_deserialize(self.julia_state_stream_),
+ jl_deserialize(stream),
)
@property
@@ -1522,6 +1729,39 @@ def raw_julia_state_(self):
def expression_spec_(self):
return self.expression_spec or ExpressionSpec()
+ def _has_fitted_type_spec(self) -> bool:
+ return hasattr(self, "_type_spec_runtime_definition_")
+
+ def _supports_export(self, format: str) -> bool:
+ return not self._has_fitted_type_spec() and bool(
+ getattr(self.expression_spec_, f"supports_{format}")
+ )
+
+ def _has_julia_backed_equations(self) -> bool:
+ """Whether `equations_` holds live Julia objects that cannot be pickled."""
+ return self._has_fitted_type_spec() or isinstance(
+ self.expression_spec_, TemplateExpressionSpec
+ )
+
+ def _define_julia_expression_types(self) -> None:
+ """Define the Julia types needed to deserialize `julia_state_stream_`."""
+ if self._has_fitted_type_spec():
+ self._load_type_spec_runtime()
+ elif isinstance(self.expression_spec_, TemplateExpressionSpec):
+ self.expression_spec_.julia_expression_spec()
+
+ def _restore_julia_backed_columns(self) -> None:
+ """Rebuild the Julia-backed equation columns dropped by pickling."""
+ if not self._has_julia_backed_equations():
+ return
+ equations = getattr(self, "equations_", None)
+ frames = equations if isinstance(equations, list) else [equations]
+ if any(
+ isinstance(frame, pd.DataFrame) and "lambda_format" not in frame.columns
+ for frame in frames
+ ):
+ self.refresh()
+
def get_best(
self, index: int | list[int] | None = None
) -> pd.Series | list[pd.Series]:
@@ -1547,6 +1787,7 @@ def get_best(
Raised when an invalid model selection strategy is provided.
"""
check_is_fitted(self, attributes=["equations_"])
+ self._restore_julia_backed_columns()
if index is not None:
if isinstance(self.equations_, list):
@@ -1611,6 +1852,63 @@ def _setup_equation_file(self):
def _clear_equation_file_contents(self):
self.equation_file_contents_ = None
+ def _operators_from_params(self) -> dict[int, list[str]]:
+ if self.operators is not None:
+ return {arity: values.copy() for arity, values in self.operators.items()}
+ operators = {
+ 2: (
+ self.binary_operators.copy()
+ if self.binary_operators is not None
+ else ["+", "-", "/", "*"]
+ )
+ }
+ if self.unary_operators is not None:
+ operators[1] = self.unary_operators.copy()
+ return operators
+
+ def _compile_type_spec_runtime(
+ self, operators: dict[int, list[str]]
+ ) -> _TypeSpecRuntimeDefinition:
+ return compile_type_spec_runtime_for_model(self, operators)
+
+ def _load_type_spec_runtime(
+ self,
+ *,
+ for_fit: bool = False,
+ operators: dict[int, list[str]] | None = None,
+ ) -> _TypeSpecRuntime:
+ if not for_fit:
+ check_is_fitted(self, attributes=["_type_spec_runtime_definition_"])
+ return load_type_spec_runtime(self._type_spec_runtime_definition_)
+
+ assert operators is not None
+ definition = self._compile_type_spec_runtime(operators)
+ fitted_definition = getattr(self, "_type_spec_runtime_definition_", None)
+ if (
+ self.warm_start
+ and fitted_definition is not None
+ and definition.fingerprint != fitted_definition.fingerprint
+ ):
+ raise ValueError(
+ "Cannot warm-start after changing the TypeSpec configuration. "
+ "Start a new search with `warm_start=False`."
+ )
+ runtime = load_type_spec_runtime(definition)
+ if (
+ fitted_definition is None
+ or definition.fingerprint != fitted_definition.fingerprint
+ ):
+ validate_type_spec_runtime(runtime)
+ return runtime
+
+ def _julia_expression_spec(
+ self,
+ type_spec_runtime: _TypeSpecRuntime | None = None,
+ ) -> AnyValue:
+ if type_spec_runtime is not None:
+ return type_spec_runtime.expression_spec
+ return self.expression_spec_.julia_expression_spec()
+
def _validate_and_modify_params(self) -> _DynamicallySetParams:
"""
Ensure parameters passed at initialization are valid.
@@ -1625,8 +1923,29 @@ def _validate_and_modify_params(self) -> _DynamicallySetParams:
values. For example, default parameters are set here
when a parameter is left set to `None`.
"""
- # Immutable parameter validation
- # Ensure instance parameters are allowable values:
+ if (
+ self.warm_start
+ and getattr(self, "julia_state_stream_", None) is not None
+ and (self.type_spec is not None) != self._has_fitted_type_spec()
+ ):
+ raise ValueError(
+ "Cannot warm-start after enabling or disabling TypeSpec. "
+ "Start a new search with `warm_start=False`."
+ )
+
+ if self.type_spec is not None:
+ validate_type_spec_model_configuration(self)
+ legacy_mutation_weights_used = any(
+ getattr(self, parameter) is not None
+ for parameter in _LEGACY_MUTATION_PARAMETERS
+ )
+ if legacy_mutation_weights_used and (
+ self.default_mutations is not None or self.mutations is not None
+ ):
+ raise ValueError(
+ "Cannot combine legacy `weight_*` parameters with "
+ "`default_mutations` or `mutations`."
+ )
# Validate operators vs binary_operators/unary_operators mutual exclusion
if self.operators is not None:
@@ -1660,7 +1979,7 @@ def _validate_and_modify_params(self) -> _DynamicallySetParams:
)
param_container = _DynamicallySetParams(
- operators={2: ["+", "*", "-", "/"]},
+ operators={2: ["+", "-", "/", "*"]},
maxdepth=self.maxsize,
constraints={},
batch_size=None,
@@ -1670,17 +1989,7 @@ def _validate_and_modify_params(self) -> _DynamicallySetParams:
)
# Convert binary_operators/unary_operators to operators format if needed
- if self.operators is None:
- # Build operators dict from binary_operators and unary_operators
- operators_dict = {}
- if self.binary_operators is not None:
- operators_dict[2] = self.binary_operators.copy()
- else:
- # Keep default binary operators
- operators_dict[2] = ["+", "*", "-", "/"]
- if self.unary_operators is not None:
- operators_dict[1] = self.unary_operators.copy()
- param_container.operators = operators_dict
+ param_container.operators = self._operators_from_params()
for param_name in map(lambda x: x.name, fields(_DynamicallySetParams)):
user_param_value = getattr(self, param_name)
@@ -1765,6 +2074,31 @@ def _validate_and_set_fit_params(
Validated units for `y`.
"""
+ if (
+ complexity_of_variables is not None
+ and self.complexity_of_variables is not None
+ ):
+ raise ValueError(
+ "You cannot set `complexity_of_variables` at both `fit` and `__init__`. "
+ "Pass it at `__init__` to set it to global default, OR use `fit` to set it for "
+ "each variable individually."
+ )
+ elif complexity_of_variables is None:
+ complexity_of_variables = self.complexity_of_variables
+
+ if self.type_spec is not None:
+ return prepare_type_spec_fit_data(
+ self,
+ X,
+ y,
+ Xresampled,
+ weights,
+ variable_names,
+ complexity_of_variables,
+ X_units,
+ y_units,
+ )
+
if isinstance(X, pd.DataFrame):
if variable_names:
variable_names = None
@@ -1789,22 +2123,6 @@ def _validate_and_set_fit_params(
"Please use valid names instead."
)
- if (
- complexity_of_variables is not None
- and self.complexity_of_variables is not None
- ):
- raise ValueError(
- "You cannot set `complexity_of_variables` at both `fit` and `__init__`. "
- "Pass it at `__init__` to set it to global default, OR use `fit` to set it for "
- "each variable individually."
- )
- elif complexity_of_variables is not None:
- complexity_of_variables = complexity_of_variables
- elif self.complexity_of_variables is not None:
- complexity_of_variables = self.complexity_of_variables
- else:
- complexity_of_variables = None
-
# Data validation and feature name fetching via sklearn
# This method sets the n_features_in_ attribute
if Xresampled is not None:
@@ -1993,8 +2311,10 @@ def _run(
y: ndarray,
runtime_params: _DynamicallySetParams,
weights: ndarray | None,
- category: ndarray | None,
seed: int,
+ type_spec_runtime: _TypeSpecRuntime | None,
+ parallelism: Literal["serial", "multithreading", "multiprocessing"],
+ numprocs: int | None,
):
"""
Run the symbolic regression fitting process on the julia backend.
@@ -2012,12 +2332,14 @@ def _run(
Weight array of the same shape as `y`.
Each element is how to weight the mean-square-error loss
for that particular element of y.
- category : ndarray | None
- If `expression_spec` is a `ParametricExpressionSpec`, then this
- argument should be a list of integers representing the category
- of each sample in `X`.
seed : int
Random seed for julia backend process.
+ type_spec_runtime : _TypeSpecRuntime | None
+ Loaded TypeSpec runtime, or `None` for numeric searches.
+ parallelism : {"serial", "multithreading", "multiprocessing"}
+ Effective Julia backend execution mode.
+ numprocs : int | None
+ Effective number of Julia worker processes.
Returns
-------
@@ -2033,9 +2355,35 @@ def _run(
# every new instance of PySRRegressor
global ALREADY_RAN
- # These are the parameters which may be modified from the ones
- # specified in init, so we define them here locally:
- operators = runtime_params.operators
+ definition_module = (
+ jl.Main
+ if type_spec_runtime is None
+ else type_spec_runtime.configuration_module
+ )
+ type_spec_operator_functions: dict[int, tuple[AnyValue, ...]] | None = None
+ if type_spec_runtime is None:
+ supports_sympy = self.expression_spec_.supports_sympy
+ operators, custom_loss, custom_full_objective, custom_loss_expression = (
+ _create_julia_operators_and_loss_functions(
+ operators=runtime_params.operators,
+ extra_sympy_mappings=self.extra_sympy_mappings,
+ supports_sympy=supports_sympy,
+ elementwise_loss=self.elementwise_loss,
+ loss_function=self.loss_function,
+ loss_function_expression=self.loss_function_expression,
+ )
+ )
+ else:
+ type_spec_operator_functions = type_spec_runtime.operator_functions
+ operators = {
+ arity: list(names)
+ for arity, names in type_spec_runtime.operator_names.items()
+ }
+ custom_loss = type_spec_runtime.elementwise_loss
+ custom_full_objective = type_spec_runtime.loss_function
+ custom_loss_expression = type_spec_runtime.loss_function_expression
+ value_type = None if type_spec_runtime is None else type_spec_runtime.value_type
+ loss_type = None
constraints = runtime_params.constraints
nested_constraints = self.nested_constraints
@@ -2047,10 +2395,6 @@ def _run(
if not ALREADY_RAN and runtime_params.update_verbosity != 0:
pysr_logger.info("Compiling Julia backend...")
- parallelism, numprocs = _map_parallelism_params(
- self.parallelism, self.procs, getattr(self, "multithreading", None)
- )
-
if self.deterministic and parallelism != "serial":
raise ValueError(
"To ensure deterministic searches, you must set `parallelism='serial'`. "
@@ -2070,12 +2414,6 @@ def _run(
"To use cluster managers, you must set `parallelism='multiprocessing'`."
)
- # TODO(mcranmer): These functions should be part of this class.
- operators = _maybe_create_inline_operators(
- operators=operators,
- extra_sympy_mappings=self.extra_sympy_mappings,
- expression_spec=self.expression_spec_,
- )
if constraints is not None:
_constraints = _process_constraints(
operators=operators,
@@ -2095,7 +2433,7 @@ def _run(
max_arity = max(operators.keys()) if operators else 2
constraints_by_arity = {arity: None for arity in range(1, max_arity + 1)}
- # Parse dict into Julia Dict for nested constraints::
+ # Parse dict into Julia Dict for nested constraints:
if nested_constraints is not None:
nested_constraints_str = "Dict("
for outer_k, outer_v in nested_constraints.items():
@@ -2104,7 +2442,9 @@ def _run(
nested_constraints_str += f"({inner_k}) => {inner_v}, "
nested_constraints_str += "), "
nested_constraints_str += ")"
- nested_constraints = jl.seval(nested_constraints_str)
+ nested_constraints = jl.Base.include_string(
+ definition_module, nested_constraints_str, "PySR nested_constraints"
+ )
# Parse dict into Julia Dict for complexities:
if complexity_of_operators is not None:
@@ -2112,47 +2452,48 @@ def _run(
for k, v in complexity_of_operators.items():
complexity_of_operators_str += f"({k}) => {v}, "
complexity_of_operators_str += ")"
- complexity_of_operators = jl.seval(complexity_of_operators_str)
+ complexity_of_operators = jl.Base.include_string(
+ definition_module,
+ complexity_of_operators_str,
+ "PySR complexity_of_operators",
+ )
# TODO: Refactor this into helper function
if isinstance(complexity_of_variables, list):
complexity_of_variables = jl_array(complexity_of_variables)
- np_dtype = self._get_precision_mapped_dtype(np.array(X))
-
- custom_loss = jl.seval(
- str(self.elementwise_loss)
- if self.elementwise_loss is not None
- else "nothing"
+ np_dtype = (
+ None
+ if type_spec_runtime is not None
+ else self._get_precision_mapped_dtype(np.array(X))
)
+
if self.elementwise_loss is not None:
- _validate_elementwise_loss(
- custom_loss,
- has_weights=weights is not None,
- probe_value=np_dtype(1.0),
- )
+ if type_spec_runtime is None:
+ assert np_dtype is not None
+ _validate_elementwise_loss(
+ custom_loss,
+ has_weights=weights is not None,
+ probe_value=np_dtype(1.0),
+ )
- custom_full_objective = jl.seval(
- str(self.loss_function) if self.loss_function is not None else "nothing"
- )
if self.loss_function is not None:
_validate_custom_full_objective(custom_full_objective)
- custom_loss_expression = jl.seval(
- str(self.loss_function_expression)
- if self.loss_function_expression is not None
- else "nothing"
- )
if self.loss_function_expression is not None:
_validate_custom_expression_objective(custom_loss_expression)
- early_stop_condition = jl.seval(
- str(self.early_stop_condition)
- if self.early_stop_condition is not None
- else "nothing"
+ early_stop_condition = (
+ jl.seval(
+ str(self.early_stop_condition)
+ if self.early_stop_condition is not None
+ else "nothing"
+ )
+ if type_spec_runtime is None
+ else type_spec_runtime.early_stop_condition
)
- input_stream = jl.seval(self.input_stream)
+ input_stream = jl.seval(_resolve_input_stream(self.input_stream))
load_required_packages(
turbo=self.turbo,
@@ -2160,14 +2501,14 @@ def _run(
autodiff_backend=self.autodiff_backend,
cluster_manager=cluster_manager,
logger_spec=self.logger_spec,
+ use_tracing=self.use_tracing,
)
if cluster_manager is not None:
active_project = jl.seval("Base.active_project()")
if isinstance(active_project, str) and len(active_project) > 0:
- # `ClusterManagers.addprocs_slurm` launches new Julia workers via `srun`.
- # The project (environment) is propagated via `JULIA_PROJECT` rather
- # than a `--project=...` flag, so ensure it is set.
+ # Some distributed worker launchers propagate the project (environment)
+ # via `JULIA_PROJECT` rather than a `--project=...` flag.
os.environ.setdefault(
"JULIA_PROJECT", str(Path(active_project).resolve().parent)
)
@@ -2178,19 +2519,27 @@ def _run(
else:
autodiff_backend = None
- mutation_weights = SymbolicRegression.MutationWeights(
- mutate_constant=self.weight_mutate_constant,
- mutate_operator=self.weight_mutate_operator,
- mutate_feature=self.weight_mutate_feature,
- swap_operands=self.weight_swap_operands,
- rotate_tree=self.weight_rotate_tree,
- add_node=self.weight_add_node,
- insert_node=self.weight_insert_node,
- delete_node=self.weight_delete_node,
- simplify=self.weight_simplify,
- randomize=self.weight_randomize,
- do_nothing=self.weight_do_nothing,
- optimize=self.weight_optimize,
+ legacy_mutation_weights = {
+ parameter.removeprefix("weight_"): getattr(self, parameter)
+ for parameter in _LEGACY_MUTATION_PARAMETERS
+ if getattr(self, parameter) is not None
+ }
+ mutation_weights = (
+ jl_named_tuple(legacy_mutation_weights) if legacy_mutation_weights else None
+ )
+ mutations = (
+ None if self.mutations is None else convert_mutations(self.mutations)
+ )
+ default_mutations = (
+ None
+ if self.default_mutations is None
+ else convert_mutations(self.default_mutations)
+ )
+ plugins = jl_array([plugin.julia_plugin() for plugin in (self.plugins or [])])
+ default_plugins = (
+ None
+ if self.default_plugins is None
+ else jl_array([plugin.julia_plugin() for plugin in self.default_plugins])
)
# Convert operators dict to Julia format and create OperatorEnum
@@ -2201,8 +2550,12 @@ def _run(
for arity in range(1, max_arity + 1):
if arity in operators:
jl_op_list = []
- for op in operators[arity]:
- jl_op = jl.seval(op)
+ for op_index, op in enumerate(operators[arity]):
+ if type_spec_runtime is None:
+ jl_op = jl.seval(op)
+ else:
+ assert type_spec_operator_functions is not None
+ jl_op = type_spec_operator_functions[arity][op_index]
if not jl_is_function(jl_op):
raise ValueError(
f"When building operators for arity {arity}, `'{op}'` did not return a Julia function"
@@ -2213,8 +2566,17 @@ def _run(
# Empty tuple for missing arities
jl_operators_dict[arity] = ()
+ if type_spec_runtime is not None:
+ loss_type = validate_type_spec_options(
+ type_spec_runtime,
+ jl_operators_dict,
+ custom_loss,
+ )
+
complexity_mapping = (
- jl.seval(self.complexity_mapping) if self.complexity_mapping else None
+ (jl.seval(self.complexity_mapping) if self.complexity_mapping else None)
+ if type_spec_runtime is None
+ else type_spec_runtime.complexity_mapping
)
if hasattr(self, "logger_") and self.logger_ is not None and self.warm_start:
@@ -2224,23 +2586,19 @@ def _run(
self.logger_ = logger
- # Use Julia function to create OperatorEnum from Dict{Int,Tuple}
- create_operator_enum = jl.seval(
- "__sr_make_op_enum(ops_dict) = OperatorEnum([k => v for (k, v) in ops_dict]...)"
+ jl_operator_enum = SymbolicRegression.OperatorEnum(
+ tuple(jl_operators_dict.values())
)
- jl_operator_enum = create_operator_enum(jl_operators_dict)
# Build constraints dict with same structure
jl_constraints_dict = None
if any(c for c in constraints_by_arity.values() if c is not None):
- constraints_pairs = []
+ jl_constraints_dict = jl.Dict[jl.Int, jl.Vector]()
for arity in range(1, max_arity + 1):
if constraints_by_arity[arity] is not None:
- constraints_pairs.append(
- jl.Pair(arity, jl_array(constraints_by_arity[arity]))
- )
- if constraints_pairs:
- jl_constraints_dict = jl.Dict(constraints_pairs)
+ jl_constraints_dict[arity] = jl_array(constraints_by_arity[arity])
+
+ expression_spec = self._julia_expression_spec(type_spec_runtime)
# Call to Julia backend.
# See https://github.com/astroautomata/SymbolicRegression.jl/blob/master/src/OptionsStruct.jl
@@ -2251,7 +2609,7 @@ def _run(
complexity_of_constants=self.complexity_of_constants,
complexity_of_variables=complexity_of_variables,
complexity_mapping=complexity_mapping,
- expression_spec=self.expression_spec_.julia_expression_spec(),
+ expression_spec=expression_spec,
nested_constraints=nested_constraints,
elementwise_loss=custom_loss,
loss_function=custom_full_objective,
@@ -2260,16 +2618,17 @@ def _run(
maxsize=int(self.maxsize),
output_directory=_escape_filename(self.output_directory_),
npopulations=int(self.populations),
- # Determine actual batching based on "auto" mode
- batching=(self.batching if self.batching != "auto" else len(X) > 1000),
- batch_size=int(
- _get_batch_size(len(X), runtime_params.batch_size)
- if (
- self.batching == True or (self.batching == "auto" and len(X) > 1000)
- )
- else len(X)
+ batching=(
+ jl.Symbol(self.batching)
+ if isinstance(self.batching, str)
+ else self.batching
),
+ batch_size=runtime_params.batch_size,
mutation_weights=mutation_weights,
+ mutations=mutations,
+ default_mutations=default_mutations,
+ plugins=plugins,
+ default_plugins=default_plugins,
tournament_selection_p=self.tournament_selection_p,
tournament_selection_n=self.tournament_selection_n,
# These have the same name:
@@ -2310,39 +2669,43 @@ def _run(
skip_mutation_failures=self.skip_mutation_failures,
max_evals=self.max_evals,
input_stream=input_stream,
+ use_tracing=self.use_tracing,
+ tracing_file=str(self.tracing_file),
early_stop_condition=early_stop_condition,
seed=seed,
deterministic=self.deterministic,
define_helper_functions=False,
)
- self.julia_options_stream_ = jl_serialize(options)
+ serialized_options = jl_serialize(options)
+ saved_state = (
+ jl_deserialize(self.julia_state_stream_)
+ if self.warm_start and self.julia_state_stream_ is not None
+ else None
+ )
+ if self.warm_start and self.julia_options_stream_ is not None:
+ SymbolicRegression.CoreModule.check_warm_start_compatibility(
+ jl_deserialize(self.julia_options_stream_), options
+ )
# Convert data to desired precision
# This converts the data into a Julia array:
- jl_X = jl_array(np.array(X, dtype=np_dtype).T)
- if len(y.shape) == 1:
- jl_y = jl_array(np.array(y, dtype=np_dtype))
+ if type_spec_runtime is not None:
+ jl_X = type_spec_to_julia_array(type_spec_runtime, X, transpose=True)
+ jl_y = type_spec_to_julia_array(type_spec_runtime, y)
else:
- jl_y = jl_array(np.array(y, dtype=np_dtype).T)
+ jl_X = jl_numpy_array(np.array(X, dtype=np_dtype).T)
+ numeric_y = np.array(y, dtype=np_dtype)
+ jl_y = jl_numpy_array(numeric_y.T if numeric_y.ndim > 1 else numeric_y)
if weights is not None:
if len(weights.shape) == 1:
- jl_weights = jl_array(np.array(weights, dtype=np_dtype))
+ jl_weights = jl_numpy_array(np.array(weights, dtype=np_dtype))
else:
- jl_weights = jl_array(np.array(weights, dtype=np_dtype).T)
+ jl_weights = jl_numpy_array(np.array(weights, dtype=np_dtype).T)
else:
jl_weights = None
- if category is not None:
- offset_for_julia_indexing = 1
- jl_category = jl_array(
- (category + offset_for_julia_indexing).astype(np.int64)
- )
- jl_extra = jl.seval("NamedTuple{(:class,)}")((jl_category,))
- else:
- jl_extra = jl.NamedTuple()
-
if len(y.shape) > 1:
# We set these manually so that they respect Python's 0 indexing
# (by default Julia will use y1, y2...)
@@ -2352,7 +2715,13 @@ def _run(
else:
jl_y_variable_names = None
- jl_guesses = _prepare_guesses_for_julia(self.guesses, self.nout_)
+ jl_guesses = _prepare_guesses_for_julia(
+ self.guesses,
+ self.nout_,
+ expression_spec=self.expression_spec_,
+ np_dtype=np_dtype,
+ type_spec_runtime=type_spec_runtime,
+ )
# Convert worker_imports to Julia symbols
jl_worker_imports = (
@@ -2360,55 +2729,73 @@ def _run(
if self.worker_imports is not None
else None
)
-
- out = SymbolicRegression.equation_search(
- jl_X,
- jl_y,
- weights=jl_weights,
- extra=jl_extra,
- niterations=int(self.niterations),
- variable_names=jl_array([str(v) for v in self.feature_names_in_]),
- display_variable_names=jl_array(
- [str(v) for v in self.display_feature_names_in_]
- ),
- y_variable_names=jl_y_variable_names,
- X_units=jl_array(self.X_units_),
- y_units=(
- jl_array(self.y_units_)
- if isinstance(self.y_units_, list)
- else self.y_units_
- ),
- options=options,
- guesses=jl_guesses,
- numprocs=numprocs,
- parallelism=parallelism,
- saved_state=self.julia_state_,
- return_state=True,
- run_id=self.run_id_,
- addprocs_function=cluster_manager,
- heap_size_hint_in_bytes=self.heap_size_hint_in_bytes,
- worker_timeout=self.worker_timeout,
- worker_imports=jl_worker_imports,
- progress=runtime_params.progress
- and self.verbosity > 0
- and len(y.shape) == 1,
- verbosity=int(self.verbosity),
- logger=logger,
- )
+ if type_spec_runtime is not None and parallelism == "multiprocessing":
+ definition = type_spec_runtime.definition
+ assert isinstance(definition, _TypeSpecRuntimeDefinition)
+ addprocs_function = create_type_spec_addprocs_function(
+ definition,
+ cluster_manager,
+ jl_worker_imports,
+ )
+ else:
+ addprocs_function = cluster_manager
+ with _external_stop_signal_context(self) as external_stop:
+ out = SymbolicRegression.equation_search._jl_call_nogil(
+ jl_X,
+ jl_y,
+ weights=jl_weights,
+ niterations=int(self.niterations),
+ variable_names=jl_array([str(v) for v in self.feature_names_in_]),
+ display_variable_names=jl_array(
+ [str(v) for v in self.display_feature_names_in_]
+ ),
+ y_variable_names=jl_y_variable_names,
+ X_units=jl_array(self.X_units_),
+ y_units=(
+ jl_array(self.y_units_)
+ if isinstance(self.y_units_, list)
+ else self.y_units_
+ ),
+ options=options,
+ guesses=jl_guesses,
+ numprocs=numprocs,
+ parallelism=parallelism,
+ saved_state=saved_state,
+ return_state=True,
+ run_id=self.run_id_,
+ addprocs_function=addprocs_function,
+ heap_size_hint_in_bytes=self.heap_size_hint_in_bytes,
+ worker_timeout=self.worker_timeout,
+ worker_imports=jl_worker_imports,
+ progress=runtime_params.progress
+ and self.verbosity > 0
+ and len(y.shape) == 1,
+ verbosity=int(self.verbosity),
+ logger=logger,
+ external_stop=external_stop,
+ **({"loss_type": loss_type} if loss_type is not None else {}),
+ )
if self.logger_spec is not None:
self.logger_spec.write_hparams(logger, self.get_params())
if not self.warm_start:
self.logger_spec.close(logger)
- self.julia_state_stream_ = jl_serialize(out)
+ serialized_state = jl_serialize(out)
+ equations = self.get_hof(out, type_spec_runtime=type_spec_runtime)
+ if type_spec_runtime is not None:
+ assert isinstance(type_spec_runtime.definition, _TypeSpecRuntimeDefinition)
+ self._type_spec_runtime_definition_ = type_spec_runtime.definition
+ self.julia_options_stream_ = serialized_options
+ self.julia_state_stream_ = serialized_state
# Set attributes
- self.equations_ = self.get_hof(out)
+ self.equations_ = equations
ALREADY_RAN = True
return self
+ @_rollback_failed_warm_start
def fit(
self,
X,
@@ -2420,7 +2807,6 @@ def fit(
complexity_of_variables: int | float | list[int | float] | None = None,
X_units: ArrayLike[str] | None = None,
y_units: str | ArrayLike[str] | None = None,
- category: ndarray | None = None,
) -> "PySRRegressor":
"""
Search for equations to fit the dataset and store them in `self.equations_`.
@@ -2451,17 +2837,12 @@ def fit(
X_units : list[str]
A list of units for each variable in `X`. Each unit should be
a string representing a Julia expression. See DynamicQuantities.jl
- https://symbolicml.org/DynamicQuantities.jl/dev/units/ for more
+ https://ai.damtp.cam.ac.uk/dynamicquantities/dev/units for more
information.
y_units : str | list[str]
Similar to `X_units`, but as a unit for the target variable, `y`.
If `y` is a matrix, a list of units should be passed. If `X_units`
is given but `y_units` is not, then `y_units` will be arbitrary.
- category : list[int]
- If `expression_spec` is a `ParametricExpressionSpec`, then this
- argument should be a list of integers representing the category
- of each sample.
-
Returns
-------
self : object
@@ -2486,19 +2867,19 @@ def fit(
self.complexity_of_variables_ = None
self.X_units_ = None
self.y_units_ = None
+ if hasattr(self, "_type_spec_runtime_definition_"):
+ del self._type_spec_runtime_definition_
self._setup_equation_file()
self._clear_equation_file_contents()
runtime_params = self._validate_and_modify_params()
+ parallelism, numprocs = _map_parallelism_params(
+ self.parallelism,
+ self.procs,
+ getattr(self, "multithreading", None),
+ )
- if category is not None:
- assert Xresampled is None
-
- if isinstance(self.expression_spec, ParametricExpressionSpec):
- assert category is not None
-
- # TODO: Put `category` here
(
X,
y,
@@ -2519,22 +2900,9 @@ def fit(
y_units,
)
- if X.shape[0] > 50000:
- warnings.warn(
- "You are using a dataset with more than 50,000 datapoints. "
- "Symbolic regression rarely benefits from this many points - consider "
- "subsampling to 10,000 points or fewer. If you have high noise, "
- "denoise the data first rather than using more points."
- )
-
random_state = check_random_state(self.random_state) # For np random
seed = cast(int, random_state.randint(0, 2**31 - 1)) # For julia random
- if isinstance(self.expression_spec, ParametricExpressionSpec):
- parametric_expression_deprecation_warning(
- self.expression_spec.max_parameters, variable_names
- )
-
# Pre transformations (feature selection and denoising)
X, y, variable_names, complexity_of_variables, X_units, y_units = (
self._pre_transform_training_data(
@@ -2562,15 +2930,32 @@ def fit(
y,
X_units,
y_units,
+ self.type_spec is None and self._supports_export("sympy"),
)
- # Initially, just save model parameters, so that
- # it can be loaded from an early exit:
- if not self.temp_equation_file:
+ type_spec_runtime = (
+ self._load_type_spec_runtime(
+ for_fit=True,
+ operators=runtime_params.operators,
+ )
+ if self.type_spec is not None
+ else None
+ )
+ if not self.temp_equation_file and not (
+ self.warm_start and self.julia_state_stream_ is not None
+ ):
self._checkpoint()
- # Perform the search:
- self._run(X, y, runtime_params, weights=weights, seed=seed, category=category)
+ self._run(
+ X,
+ y,
+ runtime_params,
+ weights=weights,
+ seed=seed,
+ type_spec_runtime=type_spec_runtime,
+ parallelism=parallelism,
+ numprocs=numprocs,
+ )
# Then, after fit, we save again, so the pickle file contains
# the equations:
@@ -2603,8 +2988,6 @@ def predict(
self,
X,
index: int | list[int] | None = None,
- *,
- category: ndarray | None = None,
) -> ndarray:
"""
Predict y from input X using the equation chosen by `model_selection`.
@@ -2621,11 +3004,6 @@ def predict(
particular row of `self.equations_`, you may specify the index here.
For multiple output equations, you must pass a list of indices
in the same order.
- category : ndarray | None
- If `expression_spec` is a `ParametricExpressionSpec`, then this
- argument should be a list of integers representing the category
- of each sample in `X`.
-
Returns
-------
y_predicted : ndarray of shape (n_samples, nout_)
@@ -2640,70 +3018,45 @@ def predict(
check_is_fitted(
self, attributes=["selection_mask_", "feature_names_in_", "nout_"]
)
+ has_type_spec = self._has_fitted_type_spec()
best_equation = self.get_best(index=index)
- # When X is an numpy array or a pandas dataframe with a RangeIndex,
- # the self.feature_names_in_ generated during fit, for the same X,
- # will cause a warning to be thrown during _validate_data.
- # To avoid this, convert X to a dataframe, apply the selection mask,
- # and then set the column/feature_names of X to be equal to those
- # generated during fit.
- if not isinstance(X, pd.DataFrame):
- X = check_array(X)
- X = pd.DataFrame(X)
- if isinstance(X.columns, pd.RangeIndex):
- if self.selection_mask_ is not None:
- # RangeIndex enforces column order allowing columns to
- # be correctly filtered with self.selection_mask_
- X = X[X.columns[self.selection_mask_]]
- X.columns = self.feature_names_in_
-
- # During fit, we replace spaces in DataFrame column names with
- # underscores. Apply the same normalization here.
- cols_str = X.columns.astype(str)
- if cols_str.str.contains(" ").any():
- X = X.copy()
- X.columns = cols_str.str.replace(" ", "_")
- warnings.warn(
- "Spaces in DataFrame column names are not supported. "
- "Spaces have been replaced with underscores. \n"
- "Please rename the columns to valid names."
- )
-
- # Without feature information, CallableEquation/lambda_format equations
- # require that the column order of X matches that of the X used during
- # the fitting process. _validate_data removes this feature information
- # when it converts the dataframe to an np array. Thus, to ensure feature
- # order is preserved after conversion, the dataframe columns must be
- # reordered/reindexed to match those of the transformed (denoised and
- # feature selected) X in fit.
- X = X.reindex(columns=self.feature_names_in_)
- X = self._validate_data_X(X)
- if self.expression_spec_.evaluates_in_julia:
- # Julia wants the right dtype
- X = X.astype(self._get_precision_mapped_dtype(X))
-
- if category is not None:
- offset_for_julia_indexing = 1
- args: tuple = (
- jl_array((category + offset_for_julia_indexing).astype(np.int64)),
- )
+ if has_type_spec:
+ X = prepare_type_spec_prediction_data(self, X)
else:
- args = ()
+ if not isinstance(X, pd.DataFrame):
+ X = pd.DataFrame(check_array(X))
+ if isinstance(X.columns, pd.RangeIndex):
+ if self.selection_mask_ is not None:
+ X = X[X.columns[self.selection_mask_]]
+ X.columns = self.feature_names_in_
+
+ columns = X.columns.astype(str)
+ if columns.str.contains(" ").any():
+ X = X.copy()
+ X.columns = columns.str.replace(" ", "_")
+ warnings.warn(
+ "Spaces in DataFrame column names are not supported. "
+ "Spaces have been replaced with underscores. \n"
+ "Please rename the columns to valid names."
+ )
+ X = X.reindex(columns=self.feature_names_in_)
+ X = self._validate_data_X(X)
+ if self.expression_spec_.evaluates_in_julia:
+ X = X.astype(self._get_precision_mapped_dtype(X))
try:
if isinstance(best_equation, list):
assert self.nout_ > 1
return np.stack(
- [
- cast(ndarray, eq["lambda_format"](X, *args))
- for eq in best_equation
- ],
+ [cast(ndarray, eq["lambda_format"](X)) for eq in best_equation],
axis=1,
)
else:
- return cast(ndarray, best_equation["lambda_format"](X, *args))
+ return cast(ndarray, best_equation["lambda_format"](X))
except Exception as error:
+ if has_type_spec:
+ raise
raise ValueError(
"Failed to evaluate the expression. "
"If you are using a custom operator, make sure to define it in `extra_sympy_mappings`, "
@@ -2712,6 +3065,14 @@ def predict(
"You can then run `model.refresh()` to re-load the expressions."
) from error
+ def score(self, X, y, sample_weight=None):
+ if self._has_fitted_type_spec():
+ raise NotImplementedError(
+ "The R^2 `score` is not defined for models using a `type_spec`. "
+ "Evaluate predictions with a metric suited to the value type."
+ )
+ return super().score(X, y, sample_weight=sample_weight)
+
def sympy(self, index: int | list[int] | None = None):
"""
Return sympy representation of the equation(s) chosen by `model_selection`.
@@ -2730,7 +3091,7 @@ def sympy(self, index: int | list[int] | None = None):
best_equation : str, list[str] of length nout_
SymPy representation of the best equation.
"""
- if not self.expression_spec_.supports_sympy:
+ if not self._supports_export("sympy"):
raise ValueError(
f"`expression_spec={self.expression_spec_}` does not support sympy export."
)
@@ -2766,7 +3127,7 @@ def latex(
best_equation : str or list[str] of length nout_
LaTeX expression of the best equation.
"""
- if not self.expression_spec_.supports_latex:
+ if not self._supports_export("latex"):
raise ValueError(
f"`expression_spec={self.expression_spec_}` does not support latex export."
)
@@ -2803,7 +3164,7 @@ def jax(self, index=None):
Dictionary of callable jax function in "callable" key,
and jax array of parameters as "parameters" key.
"""
- if not self.expression_spec_.supports_jax:
+ if not self._supports_export("jax"):
raise ValueError(
f"`expression_spec={self.expression_spec_}` does not support jax export."
)
@@ -2839,7 +3200,7 @@ def pytorch(self, index=None):
best_equation : torch.nn.Module
PyTorch module representing the expression.
"""
- if not self.expression_spec_.supports_torch:
+ if not self._supports_export("torch"):
raise ValueError(
f"`expression_spec={self.expression_spec_}` does not support torch export."
)
@@ -2903,7 +3264,12 @@ def _postprocess_dataframe(self, df: pd.DataFrame) -> pd.DataFrame:
return df
- def get_hof(self, search_output=None) -> pd.DataFrame | list[pd.DataFrame]:
+ def get_hof(
+ self,
+ search_output=None,
+ *,
+ type_spec_runtime: _TypeSpecRuntime | None = None,
+ ) -> pd.DataFrame | list[pd.DataFrame]:
"""Get the equations from a hall of fame file or search output.
If no arguments entered, the ones used
@@ -2927,6 +3293,25 @@ def get_hof(self, search_output=None) -> pd.DataFrame | list[pd.DataFrame]:
self.equation_file_contents_ = self._read_equation_file()
_validate_export_mappings(self.extra_jax_mappings, self.extra_torch_mappings)
+ if type_spec_runtime is None and self._has_fitted_type_spec():
+ type_spec_runtime = self._load_type_spec_runtime()
+ if type_spec_runtime is not None and search_output is None:
+ search_output = jl_deserialize(self.julia_state_stream_)
+
+ def create_exports(output: pd.DataFrame, output_index: int):
+ if type_spec_runtime is not None:
+ return create_type_spec_exports(
+ type_spec_runtime,
+ output,
+ search_output,
+ output_index if self.nout_ > 1 else None,
+ )
+ return self.expression_spec_.create_exports(
+ self,
+ output,
+ search_output,
+ output_index if self.nout_ > 1 else None,
+ )
equation_file_contents = cast(List[pd.DataFrame], self.equation_file_contents_)
@@ -2941,9 +3326,7 @@ def get_hof(self, search_output=None) -> pd.DataFrame | list[pd.DataFrame]:
if self.loss_scale == "log"
else []
),
- self.expression_spec_.create_exports(
- self, output, search_output, i if self.nout_ > 1 else None
- ),
+ create_exports(output, i),
],
axis=1,
),
@@ -2983,7 +3366,7 @@ def latex_table(
latex_table_str : str
A string that will render a table in LaTeX of the equations.
"""
- if not self.expression_spec_.supports_latex:
+ if not self._supports_export("latex"):
raise ValueError(
f"`expression_spec={self.expression_spec_}` does not support latex export."
)
@@ -3070,21 +3453,15 @@ def calculate_scores(df: pd.DataFrame) -> pd.DataFrame:
)
-def _prepare_guesses_for_julia(guesses, nout) -> VectorValue | None:
- """Convert Python guesses to Julia format.
-
- Parameters
- ----------
- guesses : list[str] | list[list[str]] | list[dict[str, str]] | list[list[dict[str, str]]] | None
- Initial guesses for equations
- nout : int
- Number of output dimensions
-
- Returns
- -------
- jl_guesses: VectorValue | None
- Julia-compatible guesses array or None if no guesses provided
- """
+def _prepare_guesses_for_julia(
+ guesses,
+ nout,
+ *,
+ expression_spec: AbstractExpressionSpec,
+ np_dtype: type | None,
+ type_spec_runtime: _TypeSpecRuntime | None,
+) -> VectorValue | None:
+ """Convert Python guesses to Julia format."""
if guesses is None:
return None
@@ -3112,35 +3489,52 @@ def _prepare_guesses_for_julia(guesses, nout) -> VectorValue | None:
f"Number of guess lists ({len(g)}) must match number of outputs ({nout})"
)
+ template = (
+ expression_spec if isinstance(expression_spec, TemplateExpressionSpec) else None
+ )
+ parameter_names = set(template.parameters or {}) if template is not None else set()
+
julia_guesses = []
for output_guesses in g:
julia_output_guesses = []
for item in output_guesses:
- if isinstance(item, dict):
- # Convert dict to NamedTuple for template expressions
- julia_output_guesses.append(jl_named_tuple(item))
- else:
- # Keep strings as-is
+ if not isinstance(item, dict):
julia_output_guesses.append(item)
+ continue
+ if not parameter_names:
+ julia_output_guesses.append(jl_named_tuple(item))
+ continue
+
+ converted_item = {}
+ for name, value in item.items():
+ if name not in parameter_names:
+ converted_item[name] = value
+ continue
+ if not isinstance(value, (list, np.ndarray)):
+ raise ValueError(f"Template parameter '{name}' must be a 1D vector")
+ if type_spec_runtime is not None:
+ parameter_array = object_array_1d(value)
+ else:
+ try:
+ parameter_array = np.array(value, dtype=np_dtype, copy=True)
+ except (TypeError, ValueError) as error:
+ raise ValueError(
+ f"Template parameter '{name}' must be a 1D vector "
+ "compatible with the data precision"
+ ) from error
+ if parameter_array.ndim != 1:
+ raise ValueError(f"Template parameter '{name}' must be a 1D vector")
+ converted_item[name] = (
+ type_spec_to_julia_array(type_spec_runtime, parameter_array)
+ if type_spec_runtime is not None
+ else jl_array(parameter_array)
+ )
+ julia_output_guesses.append(jl_named_tuple(converted_item))
julia_guesses.append(jl_array(julia_output_guesses))
return jl_array(julia_guesses)
-def _get_batch_size(dataset_size: int, batch_size_param: int | None) -> int:
- """Calculate the actual batch size to use."""
- if batch_size_param is not None:
- return min(dataset_size, batch_size_param)
- elif dataset_size <= 1000:
- return dataset_size
- elif dataset_size < 5000:
- return 128
- elif dataset_size < 50000:
- return 256
- else:
- return 512
-
-
def _mutate_parameter(param_name: str, param_value):
if param_name == "batch_size" and param_value is not None and param_value < 1:
warnings.warn(
@@ -3154,15 +3548,29 @@ def _mutate_parameter(param_name: str, param_value):
and param_value == True
and "buffer" not in sys.stdout.__dir__()
):
- warnings.warn(
- "Note: it looks like you are running in Jupyter. "
- "The progress bar will be turned off."
- )
+ # The progress bar needs a real stdout buffer (e.g., not Jupyter's
+ # stream proxy), so fall back to plain printing silently.
return False
return param_value
+def _resolve_input_stream(input_stream: str | None) -> str:
+ """Map `None` to ``"stdin"`` on an interactive terminal, else ``"devnull"``.
+
+ Watching stdin for the ``'q'`` quit command (and advertising it in the
+ progress output) only makes sense when a user can actually type into
+ stdin; in notebooks, pipes, and CI it cannot work.
+ """
+ if input_stream is not None:
+ return input_stream
+ try:
+ interactive = sys.stdin.isatty()
+ except (AttributeError, ValueError):
+ interactive = False
+ return "stdin" if interactive else "devnull"
+
+
def _map_parallelism_params(
parallelism: Literal["serial", "multithreading", "multiprocessing"] | None,
procs: int | None,
diff --git a/pysr/test/__init__.py b/pysr/test/__init__.py
index 0997991fd..44cb81a75 100644
--- a/pysr/test/__init__.py
+++ b/pysr/test/__init__.py
@@ -1,6 +1,7 @@
from .test_autodiff import runtests as runtests_autodiff
from .test_cli import get_runtests as get_runtests_cli
from .test_dev import runtests as runtests_dev
+from .test_interrupt import runtests as runtests_interrupt
from .test_jax import runtests as runtests_jax
from .test_main import runtests
from .test_slurm import runtests as runtests_slurm
@@ -15,5 +16,6 @@
"get_runtests_cli",
"runtests_startup",
"runtests_dev",
+ "runtests_interrupt",
"runtests_slurm",
]
diff --git a/pysr/test/slurm_docker_cluster/config/cgroup.conf b/pysr/test/slurm_docker_cluster/config/cgroup.conf
index d6b278061..dbf068535 100644
--- a/pysr/test/slurm_docker_cluster/config/cgroup.conf
+++ b/pysr/test/slurm_docker_cluster/config/cgroup.conf
@@ -5,6 +5,8 @@
# See man slurm.conf and man cgroup.conf for further
# information on cgroup configuration parameters
#
+# CI containers use linuxproc/affinity and do not run systemd.
+CgroupPlugin=disabled
ConstrainCores=no
ConstrainDevices=no
ConstrainRAMSpace=no
diff --git a/pysr/test/test_cli.py b/pysr/test/test_cli.py
index f374ddbda..cdd82b785 100644
--- a/pysr/test/test_cli.py
+++ b/pysr/test/test_cli.py
@@ -1,5 +1,7 @@
+import os
import unittest
from textwrap import dedent
+from unittest.mock import patch
from click import testing as click_testing
@@ -52,8 +54,8 @@ def test_help_on_test(self):
Run parts of the PySR test suite.
- Choose from main, jax, torch, autodiff, cli, dev, startup, and slurm. You can
- give multiple tests, separated by commas.
+ Choose from main, jax, torch, autodiff, cli, dev, startup, slurm, and
+ interrupt. You can give multiple tests, separated by commas.
Options:
-k TEXT Filter expressions to select specific tests.
@@ -63,6 +65,51 @@ def test_help_on_test(self):
self.assertEqual(result.output.strip(), expected.strip())
self.assertEqual(result.exit_code, 0)
+ def test_test_shards_partition_filtered_duplicate_groups(self):
+ executed = []
+
+ class ShardTests(unittest.TestCase):
+ def test_ignored(self):
+ executed.append(self.id())
+
+ test_selected_a = test_ignored
+ test_selected_b = test_ignored
+ test_selected_c = test_ignored
+ test_selected_d = test_ignored
+ test_selected_e = test_ignored
+
+ loader = unittest.TestLoader()
+ selected = [
+ test.id()
+ for _ in range(2)
+ for test in loader.loadTestsFromTestCase(ShardTests)
+ if "selected" in test.id()
+ ]
+ shard_runs = []
+ with patch(
+ "pysr._cli.main.get_runtests_cli",
+ return_value=lambda **_: [ShardTests],
+ ):
+ for shard_index in range(3):
+ executed.clear()
+ with patch.dict(
+ os.environ,
+ {
+ "PYSR_TEST_SHARD_COUNT": "3",
+ "PYSR_TEST_SHARD_INDEX": str(shard_index),
+ },
+ ):
+ result = self.cli_runner.invoke(
+ pysr, ["test", "cli,cli", "-k", "selected"]
+ )
+ self.assertEqual(result.exit_code, 0, result.output)
+ shard_runs.append(executed.copy())
+
+ self.assertEqual(
+ shard_runs,
+ [selected[shard_index::3] for shard_index in range(3)],
+ )
+
def runtests(just_tests=False):
"""Run all tests in cliTest.py."""
tests = [TestCli]
diff --git a/pysr/test/test_dev_pysr.dockerfile b/pysr/test/test_dev_pysr.dockerfile
index 1a8c7e776..e757a89a2 100644
--- a/pysr/test/test_dev_pysr.dockerfile
+++ b/pysr/test/test_dev_pysr.dockerfile
@@ -4,7 +4,7 @@
ARG JLVERSION=1.11
ARG PYVERSION=3.12
-ARG BASE_IMAGE=bullseye
+ARG BASE_IMAGE=trixie
FROM julia:${JLVERSION}-${BASE_IMAGE} AS jl
FROM python:${PYVERSION}-${BASE_IMAGE}
diff --git a/pysr/test/test_interrupt.py b/pysr/test/test_interrupt.py
new file mode 100644
index 000000000..83144783b
--- /dev/null
+++ b/pysr/test/test_interrupt.py
@@ -0,0 +1,398 @@
+"""End-to-end tests for graceful SIGINT handling during a search."""
+
+import ctypes
+import os
+import signal
+import subprocess
+import sys
+import tempfile
+import textwrap
+import threading
+import time
+import unittest
+import warnings
+from pathlib import Path
+from types import SimpleNamespace
+from unittest import mock
+
+import pysr.interrupt as interrupt_module
+from pysr.sr import _resolve_input_stream as sr_resolve
+
+# Native sigaction checks are only available on POSIX.
+POSIX = os.name == "posix"
+
+
+class TestResolveInputStream(unittest.TestCase):
+ def test_default_uses_stdin_on_interactive_terminal(self):
+ with mock.patch.object(sys, "stdin") as stdin:
+ stdin.isatty.return_value = True
+ self.assertEqual(sr_resolve(None), "stdin")
+
+ def test_default_disables_stdin_watching_without_terminal(self):
+ with mock.patch.object(sys, "stdin") as stdin:
+ stdin.isatty.return_value = False
+ self.assertEqual(sr_resolve(None), "devnull")
+
+ def test_default_handles_missing_or_closed_stdin(self):
+ with mock.patch.object(sys, "stdin", None):
+ self.assertEqual(sr_resolve(None), "devnull")
+ with mock.patch.object(sys, "stdin") as stdin:
+ stdin.isatty.side_effect = ValueError("I/O operation on closed file")
+ self.assertEqual(sr_resolve(None), "devnull")
+
+ def test_explicit_values_pass_through(self):
+ for value in ("stdin", "devnull", "Main.my_stream"):
+ self.assertEqual(sr_resolve(value), value)
+
+
+@unittest.skipUnless(POSIX, "SIGINT signal contexts are POSIX-only")
+class TestExternalStopSignalContext(unittest.TestCase):
+ def setUp(self):
+ self.native_handler_before = self._native_handler()
+ self.python_before = signal.getsignal(signal.SIGINT)
+ self.stop_fds = []
+ backend = SimpleNamespace(ExternalStop=self._external_stop)
+ self.patches = (
+ mock.patch.dict(os.environ, {"PYTHON_JULIACALL_HANDLE_SIGNALS": "yes"}),
+ mock.patch.object(interrupt_module, "SymbolicRegression", backend),
+ )
+ for patcher in self.patches:
+ patcher.start()
+ self.addCleanup(patcher.stop)
+
+ def tearDown(self):
+ self.assertEqual(signal.getsignal(signal.SIGINT), self.python_before)
+ self.assertEqual(self._native_handler(), self.native_handler_before)
+
+ def _external_stop(self, read_fd, trigger):
+ self.stop_fds.append(read_fd)
+ return SimpleNamespace(fd=read_fd, trigger=trigger)
+
+ @staticmethod
+ def _native_handler():
+ # Compare only `sa_handler`. It is the whole of what this test
+ # asserts, and as the first member on both Linux and Darwin it needs
+ # no platform-specific struct offsets to read.
+ class Action(ctypes.Union):
+ _fields_ = [
+ ("handler", ctypes.c_void_p),
+ ("alignment", ctypes.c_longdouble),
+ ("storage", ctypes.c_ubyte * 1024),
+ ]
+
+ libc = ctypes.CDLL(None, use_errno=True)
+ libc.sigaction.argtypes = [ctypes.c_int, ctypes.c_void_p, ctypes.c_void_p]
+ libc.sigaction.restype = ctypes.c_int
+ action = Action()
+ result = libc.sigaction(signal.SIGINT, None, ctypes.byref(action))
+ if result != 0:
+ errno = ctypes.get_errno()
+ raise OSError(errno, os.strerror(errno))
+ return action.handler
+
+ @staticmethod
+ def _current_wakeup_fd():
+ current = signal.set_wakeup_fd(-1)
+ signal.set_wakeup_fd(current)
+ return current
+
+ def _assert_stop_fds_closed(self):
+ for fd in self.stop_fds:
+ with self.assertRaises(OSError):
+ os.fstat(fd)
+
+ def test_cleanup_after_normal_return_restores_existing_wakeup_fd(self):
+ wakeup_read, wakeup_write = os.pipe()
+ os.set_blocking(wakeup_write, False)
+ previous_wakeup = signal.set_wakeup_fd(wakeup_write)
+ try:
+ model = SimpleNamespace()
+ with interrupt_module._external_stop_signal_context(model) as external_stop:
+ self.assertEqual(external_stop.fd, self.stop_fds[0])
+ self.assertEqual(external_stop.trigger, signal.SIGINT)
+ self.assertFalse(os.get_blocking(external_stop.fd))
+ self.assertFalse(os.get_blocking(self._current_wakeup_fd()))
+ self.assertEqual(self._current_wakeup_fd(), wakeup_write)
+ self.assertFalse(model.interrupted_)
+ self._assert_stop_fds_closed()
+ finally:
+ signal.set_wakeup_fd(previous_wakeup)
+ os.close(wakeup_write)
+ os.close(wakeup_read)
+
+ def test_cleanup_after_search_error(self):
+ model = SimpleNamespace()
+ with self.assertRaisesRegex(RuntimeError, "search failed"):
+ with interrupt_module._external_stop_signal_context(model):
+ raise RuntimeError("search failed")
+ self._assert_stop_fds_closed()
+
+ def test_cleanup_continues_after_cleanup_failure(self):
+ real_set_wakeup_fd = signal.set_wakeup_fd
+ calls = []
+
+ def fail_after_restoring(fd):
+ result = real_set_wakeup_fd(fd)
+ calls.append(fd)
+ if len(calls) == 2:
+ raise RuntimeError("injected cleanup failure")
+ return result
+
+ model = SimpleNamespace()
+ with mock.patch.object(signal, "set_wakeup_fd", fail_after_restoring):
+ with self.assertRaisesRegex(RuntimeError, "injected cleanup failure"):
+ with interrupt_module._external_stop_signal_context(model):
+ pass
+ self._assert_stop_fds_closed()
+
+ def test_signals_disabled_does_not_arm(self):
+ model = SimpleNamespace(interrupted_=True)
+ with (
+ mock.patch.dict(os.environ, {"PYTHON_JULIACALL_HANDLE_SIGNALS": "no"}),
+ mock.patch.object(os, "pipe") as pipe,
+ ):
+ with interrupt_module._external_stop_signal_context(model) as external_stop:
+ self.assertIsNone(external_stop)
+ pipe.assert_not_called()
+ self.assertFalse(model.interrupted_)
+
+ def test_interrupted_return_sets_attribute_and_warns(self):
+ model = SimpleNamespace()
+ with warnings.catch_warnings(record=True) as caught:
+ warnings.simplefilter("always")
+ with interrupt_module._external_stop_signal_context(model):
+ handler = signal.getsignal(signal.SIGINT)
+ handler(signal.SIGINT, None)
+ self.assertTrue(model.interrupted_)
+ self.assertTrue(any("partial" in str(item.message).lower() for item in caught))
+
+
+# The child asserts partial results, restored native SIGINT disposition on
+# POSIX, and that a second fit in the same process works.
+CHILD_SCRIPT = textwrap.dedent("""
+ import ctypes
+ import os
+ import signal
+ import threading
+ import time
+ import warnings
+
+ import numpy as np
+
+ from pysr import PySRRegressor
+
+ if os.name == "posix":
+ libc = ctypes.CDLL(None)
+
+ def native_handler():
+ # Compare only `sa_handler`: the first member on both platforms, and
+ # all this check needs.
+ storage = (ctypes.c_char * 512)()
+ libc.sigaction(int(signal.SIGINT), None, storage)
+ return ctypes.cast(storage, ctypes.POINTER(ctypes.c_void_p)).contents.value
+
+ before = native_handler()
+
+ def send_interrupt():
+ if os.name == "posix":
+ os.kill(os.getpid(), signal.SIGINT)
+ elif not ctypes.windll.kernel32.GenerateConsoleCtrlEvent(0, 0):
+ raise ctypes.WinError()
+
+ rstate = np.random.RandomState(0)
+ X = rstate.randn(150, 2)
+ y = X[:, 0] * X[:, 1]
+ warmup = PySRRegressor(
+ niterations=1,
+ populations=3,
+ verbosity=0,
+ progress=False,
+ temp_equation_file=True,
+ )
+ warmup.fit(X[:20], y[:20])
+
+ model = PySRRegressor(
+ niterations=1_000_000, # only an interrupt can end this fit
+ populations=8,
+ verbosity=0,
+ progress=False,
+ temp_equation_file=True,
+ )
+ unarmed_handler = signal.getsignal(signal.SIGINT)
+
+ def interrupt_once_armed():
+ # The fit installs its own SIGINT handler once the search is armed; the
+ # extra wait lets the search run long enough to have partial results.
+ while signal.getsignal(signal.SIGINT) is unarmed_handler:
+ time.sleep(0.1)
+ time.sleep(5.0)
+ send_interrupt()
+
+ threading.Thread(target=interrupt_once_armed, daemon=True).start()
+ with warnings.catch_warnings(record=True) as caught:
+ warnings.simplefilter("always")
+ model.fit(X, y)
+ assert model.equations_ is not None
+ assert len(model.equations_) > 0, "search returned no partial results"
+ assert model.interrupted_ is True
+ assert any("partial" in str(item.message).lower() for item in caught)
+ print(f"INTERRUPTED_OK:{len(model.equations_)}", flush=True)
+
+ if os.name == "posix":
+ print(f"HANDLER_RESTORED:{native_handler() == before}", flush=True)
+
+ model2 = PySRRegressor(
+ niterations=2,
+ populations=8,
+ verbosity=0,
+ progress=False,
+ temp_equation_file=True,
+ )
+ model2.fit(X, y)
+ assert model2.equations_ is not None
+ assert model2.interrupted_ is False
+ print("SECOND_FIT_OK", flush=True)
+ """)
+
+# Wait this long after Jupyter's SEARCHING marker so the fit has armed the
+# cooperative handler before the single user interrupt.
+FIRST_SIGNAL_DELAY = 15.0
+TOTAL_TIMEOUT = 600.0
+
+
+class TestSubprocessInterrupt(unittest.TestCase):
+ def test_sigint_returns_partial_results_and_restores_state(self):
+ with tempfile.TemporaryDirectory() as tmpdir:
+ script = Path(tmpdir) / "child.py"
+ script.write_text(CHILD_SCRIPT)
+ p = subprocess.Popen(
+ [sys.executable, str(script)],
+ cwd=tmpdir,
+ stdout=subprocess.PIPE,
+ stderr=subprocess.STDOUT,
+ text=True,
+ creationflags=0 if POSIX else subprocess.CREATE_NEW_CONSOLE,
+ )
+ watchdog = threading.Timer(TOTAL_TIMEOUT, p.kill)
+ watchdog.start()
+ lines = []
+ try:
+ for line in p.stdout:
+ lines.append(line.strip())
+ p.wait()
+ finally:
+ watchdog.cancel()
+ if p.poll() is None:
+ p.kill()
+
+ output = "\n".join(lines)
+ self.assertEqual(p.returncode, 0, f"child failed:\n{output}")
+ self.assertIn("INTERRUPTED_OK:", output)
+ if POSIX:
+ self.assertIn("HANDLER_RESTORED:True", output)
+ self.assertIn("SECOND_FIT_OK", output)
+
+
+JUPYTER_FIT_CELL = textwrap.dedent("""
+ import numpy as np
+ from pysr import PySRRegressor
+
+ rstate = np.random.RandomState(0)
+ X = rstate.randn(150, 2)
+ y = X[:, 0] * X[:, 1]
+
+ # Compile the search pipeline before the timed window opens, so that the
+ # interrupt lands on a running search rather than on Julia's warm-up.
+ warmup = PySRRegressor(
+ niterations=1,
+ populations=3,
+ verbosity=0,
+ progress=False,
+ temp_equation_file=True,
+ )
+ warmup.fit(X[:20], y[:20])
+
+ model = PySRRegressor(
+ niterations=1_000_000,
+ populations=8,
+ verbosity=0,
+ progress=False,
+ temp_equation_file=True,
+ )
+ print("SEARCHING", flush=True)
+ model.fit(X, y)
+ assert model.interrupted_ is True
+ assert len(model.equations_) > 0, "search returned no partial results"
+ print("INTERRUPTED_OK", len(model.equations_), flush=True)
+ """)
+
+
+class TestJupyterInterrupt(unittest.TestCase):
+ """Drive a real ipykernel and interrupt it exactly like the Jupyter UI."""
+
+ def test_kernel_interrupt_stops_search_and_kernel_survives(self):
+ try:
+ from jupyter_client.manager import start_new_kernel
+ except ImportError:
+ self.skipTest("jupyter_client is not installed")
+
+ km, kc = start_new_kernel()
+ try:
+ msg_id = kc.execute(JUPYTER_FIT_CELL)
+ self._await_stream(kc, "SEARCHING", timeout=TOTAL_TIMEOUT)
+
+ # Interrupt once, exactly like one click on Jupyter's interrupt action.
+ time.sleep(FIRST_SIGNAL_DELAY)
+ km.interrupt_kernel()
+ reply = self._await_reply(kc, msg_id, timeout=TOTAL_TIMEOUT)
+ assert reply is not None, "kernel never returned from fit"
+ # The fit returns normally with partial results: no error status.
+ self.assertEqual(reply["content"]["status"], "ok")
+
+ # The kernel must remain fully usable afterwards.
+ reply = kc.execute_interactive(
+ "print('ALIVE', len(model.equations_), flush=True)",
+ timeout=60,
+ )
+ self.assertEqual(reply["content"]["status"], "ok")
+ finally:
+ kc.stop_channels()
+ km.shutdown_kernel(now=True)
+
+ # -- helpers -----------------------------------------------------------
+
+ def _await_stream(self, kc, text, timeout):
+ deadline = time.monotonic() + timeout
+ while time.monotonic() < deadline:
+ msg = kc.get_iopub_msg(timeout=deadline - time.monotonic())
+ if msg["msg_type"] == "stream" and text in msg["content"]["text"]:
+ return
+ raise AssertionError(f"never saw {text!r} on iopub")
+
+ def _await_reply(self, kc, msg_id, timeout):
+ deadline = time.monotonic() + timeout
+ while time.monotonic() < deadline:
+ try:
+ msg = kc.get_shell_msg(timeout=max(0.1, deadline - time.monotonic()))
+ except Exception:
+ return None
+ if msg["parent_header"].get("msg_id") == msg_id:
+ return msg
+ return None
+
+
+def runtests(just_tests=False):
+ tests = [
+ TestResolveInputStream,
+ TestExternalStopSignalContext,
+ TestSubprocessInterrupt,
+ TestJupyterInterrupt,
+ ]
+ if just_tests:
+ return tests
+ suite = unittest.TestSuite()
+ loader = unittest.TestLoader()
+ for test in tests:
+ suite.addTests(loader.loadTestsFromTestCase(test))
+ runner = unittest.TextTestRunner()
+ return runner.run(suite)
diff --git a/pysr/test/test_jax.py b/pysr/test/test_jax.py
index d261afcf5..45db9f3c7 100644
--- a/pysr/test/test_jax.py
+++ b/pysr/test/test_jax.py
@@ -1,5 +1,7 @@
+import subprocess
+import sys
+import tempfile
import unittest
-from functools import partial
from pathlib import Path
import numpy as np
@@ -49,7 +51,7 @@ def test_pipeline_pandas(self):
for fname in ["hall_of_fame.csv.bak", "hall_of_fame.csv"]:
equations["Complexity Loss Equation".split(" ")].to_csv(
- Path(model.output_directory_) / model.run_id_ / fname
+ Path(model.output_directory_) / model.run_id_ / fname, index=False
)
model.refresh(run_directory=str(Path(model.output_directory_) / model.run_id_))
@@ -77,7 +79,7 @@ def test_pipeline(self):
for fname in ["hall_of_fame.csv.bak", "hall_of_fame.csv"]:
equations["Complexity Loss Equation".split(" ")].to_csv(
- Path(model.output_directory_) / model.run_id_ / fname
+ Path(model.output_directory_) / model.run_id_ / fname, index=False
)
model.refresh(run_directory=str(Path(model.output_directory_) / model.run_id_))
@@ -125,29 +127,53 @@ def cos_approx(x):
return 1 - (x**2) / 2 + (x**4) / 24 + (x**6) / 720
sp_cos_approx = sympy.Function("cos_approx")
-
y = X["k15"] ** 2 + 2 * cos_approx(X["k20"])
- model = PySRRegressor(
- progress=False,
- unary_operators=["cos_approx(x) = 1 - x^2 / 2 + x^4 / 24 + x^6 / 720"],
- select_k_features=3,
- maxsize=10,
- early_stop_condition=1e-5,
- extra_sympy_mappings={"cos_approx": sp_cos_approx},
- extra_jax_mappings={
- sp_cos_approx: "(lambda x: 1 - x**2 / 2 + x**4 / 24 + x**6 / 720)"
- },
- random_state=0,
- deterministic=True,
- parallelism="serial",
- )
- np.random.seed(0)
- model.fit(X.values, y.values)
- f, parameters = model.jax().values()
- jax_prediction = partial(f, parameters=parameters)
- jax_output = jax_prediction(X.values)
- np.testing.assert_almost_equal(y.values, jax_output, decimal=3)
+ with tempfile.TemporaryDirectory() as directory:
+ model = PySRRegressor(
+ progress=False,
+ unary_operators=["cos_approx(x) = 1 - x^2 / 2 + x^4 / 24 + x^6 / 720"],
+ select_k_features=3,
+ maxsize=10,
+ early_stop_condition=1e-5,
+ extra_sympy_mappings={"cos_approx": sp_cos_approx},
+ extra_jax_mappings={
+ sp_cos_approx: "(lambda x: 1 - x**2 / 2 + x**4 / 24 + x**6 / 720)"
+ },
+ random_state=0,
+ deterministic=True,
+ parallelism="serial",
+ output_directory=directory,
+ run_id="custom-jax-checkpoint",
+ temp_equation_file=False,
+ )
+ np.random.seed(0)
+ model.fit(X.values, y.values)
+ self.assertTrue(model.show_pickle_warnings_)
+
+ input_path = Path(directory) / "input.npy"
+ output_path = Path(directory) / "output.npy"
+ np.save(input_path, X.values)
+ run_directory = Path(directory) / "custom-jax-checkpoint"
+ code = f"""
+import numpy as np
+from pysr import PySRRegressor
+model = PySRRegressor.from_file(run_directory={str(run_directory)!r})
+f, parameters = model.jax().values()
+np.save({str(output_path)!r}, np.asarray(f(np.load({str(input_path)!r}), parameters)))
+"""
+ result = subprocess.run(
+ [sys.executable, "-c", code],
+ check=False,
+ capture_output=True,
+ text=True,
+ )
+ self.assertEqual(result.returncode, 0, result.stderr)
+ symbols = [sympy.Symbol(name) for name in model.feature_names_in_]
+ reference = sympy.lambdify(
+ symbols, model.sympy(), modules=[{"cos_approx": cos_approx}, "numpy"]
+ )(*X.values[:, model.selection_mask_].T)
+ np.testing.assert_almost_equal(reference, np.load(output_path), decimal=3)
def runtests(just_tests=False):
diff --git a/pysr/test/test_main.py b/pysr/test/test_main.py
index b703ff752..453eaa50a 100644
--- a/pysr/test/test_main.py
+++ b/pysr/test/test_main.py
@@ -1,13 +1,19 @@
import functools
import importlib
+import json
import os
import pickle as pkl
import platform
-import re
+import shutil
+import subprocess
+import sys
import tempfile
+import threading
+import time
import traceback
import unittest
import warnings
+from dataclasses import fields, is_dataclass
from pathlib import Path
from textwrap import dedent
from unittest import mock
@@ -24,7 +30,6 @@
estimator_checks_generator = functools.partial(check_estimator, generate_only=True)
from pysr import (
- ParametricExpressionSpec,
PySRRegressor,
TemplateExpressionSpec,
TensorBoardLoggerSpec,
@@ -34,11 +39,11 @@
)
from pysr.export_latex import sympy2latex
from pysr.export_sympy import pysr2sympy
-from pysr.expression_specs import parametric_expression_deprecation_warning
from pysr.feature_selection import _handle_feature_selection, run_feature_selection
-from pysr.julia_helpers import init_julia
+from pysr.julia_helpers import _load_cluster_manager, init_julia
from pysr.sr import (
_check_assertions,
+ _create_julia_operators_and_loss_functions,
_process_constraints,
_suggest_keywords,
_validate_elementwise_loss,
@@ -53,6 +58,7 @@
DEFAULT_POPULATIONS,
skip_if_beartype,
)
+from .test_type_specs import TestTypeSpecs
# Disables local saving:
os.environ["SYMBOLIC_REGRESSION_IS_TESTING"] = os.environ.get(
@@ -353,6 +359,17 @@ def test_elementwise_loss_float32_fit_accepts_strictly_typed_loss(self):
def test_validation_helpers_skip_nonfunction(self):
_validate_elementwise_loss(jl.seval("1.0"), has_weights=False)
+ def test_elementwise_loss_accepts_lossfunctions_object(self):
+ _, custom_loss, _, _ = _create_julia_operators_and_loss_functions(
+ operators={2: ["+", "-", "/", "*"]},
+ extra_sympy_mappings=None,
+ supports_sympy=True,
+ elementwise_loss="LPDistLoss{3}()",
+ loss_function=None,
+ loss_function_expression=None,
+ )
+ self.assertTrue(jl.applicable(custom_loss, 1.0, 1.0))
+
def test_validate_export_mappings_typechecks(self):
with self.assertRaises(ValueError):
_validate_export_mappings({"a": 1}, None)
@@ -621,6 +638,10 @@ def test_warm_start_set_at_init(self):
regressor = PySRRegressor(warm_start=True, max_evals=10)
regressor.fit(self.X, y)
+ regressor.set_params(binary_operators=["+", "*", "-", "/"])
+ with self.assertRaisesRegex(JuliaError, "Warm start incompatible.*operators"):
+ regressor.fit(self.X, y)
+
def test_noisy_builtin_variable_names(self):
y = self.X[:, [0, 1]] ** 2 + self.rstate.randn(self.X.shape[0], 1) * 0.05
model = PySRRegressor(
@@ -832,7 +853,9 @@ def test_template_expressions_and_custom_complexity(self):
# Create model with template that includes the missing sin operator
model = PySRRegressor(
expression_spec=TemplateExpressionSpec(
- ["f"], "sin_of_f((; f), (x, y)) = sin(f(x, y))"
+ combine="sin(f(x, y))",
+ expressions=["f"],
+ variable_names=["x", "y"],
),
binary_operators=["+", "-", "*", "/"],
unary_operators=[], # No sin operator!
@@ -879,6 +902,112 @@ def test_template_expressions_and_custom_complexity(self):
with self.assertRaises(ValueError):
model.latex_table()
+ def test_template_expressions_reload_in_fresh_process(self):
+ X = self.rstate.uniform(-1, 1, (30, 2))
+ y = np.sin(X[:, 0] + X[:, 1])
+ with tempfile.TemporaryDirectory() as directory:
+ model = PySRRegressor(
+ expression_spec=TemplateExpressionSpec(
+ combine="sin(f(x, y))",
+ expressions=["f"],
+ variable_names=["x", "y"],
+ ),
+ binary_operators=["+", "-", "*"],
+ unary_operators=[],
+ maxsize=10,
+ **{
+ **self.default_test_kwargs,
+ "temp_equation_file": False,
+ },
+ output_directory=directory,
+ run_id="template-checkpoint",
+ early_stop_condition=1e-10,
+ )
+ model.fit(X, y)
+ expected = model.predict(X)
+ run_directory = Path(directory) / "template-checkpoint"
+ code = f"""
+import json
+import numpy as np
+from pysr import PySRRegressor
+model = PySRRegressor.from_file(run_directory={str(run_directory)!r})
+X = np.array({X.tolist()!r})
+print(json.dumps({{
+ "columns": list(model.equations_.columns),
+ "prediction": model.predict(X).tolist(),
+}}))
+"""
+ result = subprocess.run(
+ [sys.executable, "-c", code],
+ check=False,
+ capture_output=True,
+ text=True,
+ )
+
+ self.assertEqual(result.returncode, 0, result.stderr)
+ payload = json.loads(result.stdout.strip().splitlines()[-1])
+ self.assertIn("lambda_format", payload["columns"])
+ np.testing.assert_allclose(payload["prediction"], expected, atol=1e-10)
+
+ loaded = pkl.loads(pkl.dumps(model))
+ self.assertIsNotNone(loaded.julia_state_)
+
+ def test_template_expression_pickle_roundtrip_in_fresh_process(self):
+ X = self.rstate.uniform(-1, 1, (30, 2))
+ y = np.sin(X[:, 0] + X[:, 1])
+ with tempfile.TemporaryDirectory() as directory:
+ model = PySRRegressor(
+ expression_spec=TemplateExpressionSpec(
+ combine="sin(f(x, y))",
+ expressions=["f"],
+ variable_names=["x", "y"],
+ ),
+ binary_operators=["+", "-", "*"],
+ unary_operators=[],
+ maxsize=10,
+ **{
+ **self.default_test_kwargs,
+ "temp_equation_file": False,
+ },
+ output_directory=directory,
+ run_id="template-pickle",
+ early_stop_condition=1e-10,
+ )
+ model.fit(X, y)
+ expected = model.predict(X)
+
+ pickle_path = Path(directory) / "template-model.pkl"
+ with open(pickle_path, "wb") as f:
+ pkl.dump(model, f)
+ # Delete the run directory: the pickle alone must be enough to
+ # rebuild the Julia-backed equation columns.
+ shutil.rmtree(Path(directory) / "template-pickle")
+
+ code = f"""
+import json
+import pickle
+import numpy as np
+
+with open({str(pickle_path)!r}, "rb") as f:
+ model = pickle.load(f)
+X = np.array({X.tolist()!r})
+prediction = model.predict(X)
+print(json.dumps({{
+ "columns": list(model.equations_.columns),
+ "prediction": prediction.tolist(),
+}}))
+"""
+ result = subprocess.run(
+ [sys.executable, "-c", code],
+ check=False,
+ capture_output=True,
+ text=True,
+ )
+ self.assertEqual(result.returncode, 0, result.stderr)
+ payload = json.loads(result.stdout.strip().splitlines()[-1])
+ self.assertIn("lambda_format", payload["columns"])
+ np.testing.assert_allclose(payload["prediction"], expected, atol=1e-10)
+
def test_template_expression_with_parameters_multiout(self):
# Create random data
X_continuous = self.rstate.uniform(-1, 1, (100, 2))
@@ -897,7 +1026,7 @@ def test_template_expression_with_parameters_multiout(self):
model = PySRRegressor(
**self.default_test_kwargs,
expression_spec=TemplateExpressionSpec(
- "p[class] * x1^2 + f(x2)",
+ combine="p[class] * x1^2 + f(x2)",
expressions=["f"],
parameters={"p": 3},
variable_names=["x1", "x2", "class"],
@@ -927,61 +1056,6 @@ def test_template_expression_with_parameters_multiout(self):
test_mse = np.mean((y_test - y_pred) ** 2)
self.assertLess(test_mse, 1e-5)
- def test_parametric_expression(self):
- # Create data with two classes
- n_points = 100
- X = self.rstate.uniform(-3, 3, (n_points, 2)) # x1, x2
- category = self.rstate.randint(0, 3, n_points) # class (0 or 1)
-
- # True parameters for each class
- P1 = [0.1, 1.5, -5.2] # phase shift for each class
- P2 = [3.2, 0.5, 1.2] # offset for each class
-
- # Ground truth: 2*cos(x2 + P1[class]) + x1^2 - P2[class]
- y = np.array(
- [
- 2 * np.cos(x2 + P1[c]) + x1**2 - P2[c]
- for x1, x2, c in zip(X[:, 0], X[:, 1], category)
- ]
- )
-
- model = PySRRegressor(
- expression_spec=ParametricExpressionSpec(max_parameters=2),
- binary_operators=["+", "*", "/", "-"],
- unary_operators=["cos", "exp"],
- maxsize=20,
- early_stop_condition="stop_if(loss, complexity) = loss < 1e-4 && complexity <= 14",
- **self.default_test_kwargs,
- )
-
- model.fit(X, y, category=category)
-
- # Test on new data points
- X_test = self.rstate.uniform(-6, 6, (10, 2))
- category_test = self.rstate.randint(0, 3, 10)
-
- y_test = np.array(
- [
- 2 * np.cos(x2 + P1[c]) + x1**2 - P2[c]
- for x1, x2, c in zip(X_test[:, 0], X_test[:, 1], category_test)
- ]
- )
-
- y_test_pred = model.predict(X_test, category=category_test)
- test_mse = np.mean((y_test - y_test_pred) ** 2)
- self.assertLess(test_mse, 1e-3)
-
- with self.assertRaises(ValueError):
- model.sympy()
- with self.assertRaises(ValueError):
- model.latex()
- with self.assertRaises(ValueError):
- model.jax()
- with self.assertRaises(ValueError):
- model.pytorch()
- with self.assertRaises(ValueError):
- model.latex_table()
-
def test_tensorboard_logger(self):
if platform.system() == "Windows":
@@ -1043,6 +1117,37 @@ def test_tensorboard_logger(self):
# Verify model still works as expected
self.assertLessEqual(model.get_best()["loss"], 1e-4)
+ def test_tracing_writes_jsonl(self):
+ """Tracing should write the search and each iteration as JSON lines."""
+ X = self.rstate.randn(10, 2)
+ y = X[:, 0]
+ kwargs = dict(
+ niterations=2,
+ populations=2,
+ population_size=20,
+ ncycles_per_iteration=DEFAULT_NCYCLES,
+ progress=False,
+ temp_equation_file=True,
+ parallelism="serial",
+ deterministic=True,
+ random_state=0,
+ )
+ with tempfile.TemporaryDirectory() as tmpdir:
+ trace_file = Path(tmpdir) / "trace.jsonl"
+
+ PySRRegressor(**kwargs, use_tracing=True, tracing_file=trace_file).fit(X, y)
+
+ records = [json.loads(line) for line in trace_file.read_text().splitlines()]
+ self.assertEqual(records[0]["record_type"], "search")
+ iterations = [r for r in records if r["record_type"] == "iteration"]
+ self.assertGreater(len(iterations), 0)
+ self.assertIn("tree", iterations[0]["members"][0])
+ self.assertIn("mutations", iterations[0])
+
+ trace_file.unlink()
+ PySRRegressor(**kwargs).fit(X, y)
+ self.assertFalse(trace_file.exists())
+
def test_negative_losses(self):
X = self.rstate.rand(100, 3) * 20.0
variance = X[:, 0] * 0.1
@@ -1191,6 +1296,29 @@ def test_template_expression_guesses(self):
# Check that a good fit was found
self.assertTrue(any(model.equations_["loss"] < 1e-10))
+ def test_template_expression_guess_parameters(self):
+ X = self.rstate.randn(100, 1)
+ y = 5.0 * X[:, 0] ** 2 + 10.0 * X[:, 0] + 0.8
+ parameter_values = np.array([5.0, 10.0, 0.8], dtype=np.float64)
+ original_values = parameter_values.copy()
+ model = PySRRegressor(
+ expression_spec=TemplateExpressionSpec(
+ expressions=["f"],
+ combine="p[1] * f(x) + p[2] * x + p[3]",
+ variable_names=["x"],
+ parameters={"p": 3, "unused": 1},
+ ),
+ guesses=[{"f": "#1 * #1", "p": parameter_values}],
+ should_optimize_constants=False,
+ precision=32,
+ **self.default_test_kwargs,
+ )
+ model.fit(X, y, variable_names=["x"])
+
+ np.testing.assert_allclose(model.predict(X), y, rtol=1e-6, atol=1e-6)
+ np.testing.assert_array_equal(parameter_values, original_values)
+ self.assertEqual(parameter_values.dtype, np.float64)
+
def test_multi_output_template_expression_guesses(self):
X = self.rstate.randn(100, 2)
Y = np.column_stack([X[:, 0] + X[:, 1], X[:, 0] - X[:, 1]])
@@ -1210,6 +1338,77 @@ def test_multi_output_template_expression_guesses(self):
for i, eqs_df in enumerate(model.equations_):
self.assertTrue(any(eqs_df["loss"] < 1e-10))
+ def test_template_parameter_guesses_replace_on_warm_start(self):
+ X = self.rstate.randn(100, 1)
+ y = 3.0 * X[:, 0] + 4.0
+ template = TemplateExpressionSpec(
+ expressions=["f"],
+ combine="p[1] * f(x) + p[2]",
+ variable_names=["x"],
+ parameters={"p": 2},
+ )
+ model = PySRRegressor(
+ expression_spec=template,
+ guesses=[{"f": "#1", "p": [1.0, 0.0]}],
+ should_optimize_constants=False,
+ **self.default_test_kwargs,
+ )
+ model.fit(X, y, variable_names=["x"])
+ model.set_params(guesses=[{"f": "#1", "p": [3.0, 4.0]}], warm_start=True)
+ model.fit(X, y, variable_names=["x"])
+
+ np.testing.assert_allclose(model.predict(X), y, rtol=1e-6, atol=1e-6)
+
+ def test_multi_output_template_parameter_guesses(self):
+ X = self.rstate.randn(100, 1)
+ Y = np.column_stack([2.0 * X[:, 0] + 1.0, -3.0 * X[:, 0] + 0.5])
+ model = PySRRegressor(
+ expression_spec=TemplateExpressionSpec(
+ expressions=["f"],
+ combine="p[1] * f(x) + p[2]",
+ variable_names=["x"],
+ parameters={"p": 2},
+ ),
+ guesses=[
+ [{"f": "#1", "p": [2.0, 1.0]}],
+ [{"f": "#1", "p": [-3.0, 0.5]}],
+ ],
+ should_optimize_constants=False,
+ **self.default_test_kwargs,
+ )
+ model.fit(X, Y, variable_names=["x"])
+
+ prediction = model.predict(X)
+ self.assertEqual(prediction.shape, Y.shape)
+ np.testing.assert_allclose(prediction, Y, rtol=1e-6, atol=1e-6)
+
+ def test_invalid_template_parameter_guesses(self):
+ X = self.rstate.randn(10, 1)
+ y = X[:, 0]
+ template = TemplateExpressionSpec(
+ expressions=["f"],
+ combine="p[1] * f(x) + p[2]",
+ variable_names=["x"],
+ parameters={"p": 2},
+ )
+ cases = [
+ ({"f": "#1", "missing": "#1", "p": [1.0, 2.0]}, JuliaError),
+ ({"p": [1.0, 2.0]}, JuliaError),
+ ({"f": ["#1"], "p": [1.0, 2.0]}, JuliaError),
+ ({"f": "#1", "p": [1.0]}, JuliaError),
+ ({"f": "#1", "p": 1.0}, ValueError),
+ ({"f": "#1", "p": [[1.0, 2.0]]}, ValueError),
+ ]
+ for guess, error in cases:
+ with self.subTest(guess=guess):
+ model = PySRRegressor(
+ expression_spec=template,
+ guesses=[guess],
+ **self.default_test_kwargs,
+ )
+ with self.assertRaises(error):
+ model.fit(X, y, variable_names=["x"])
+
def test_invalid_multi_output_format_guesses(self):
X = self.rstate.randn(100, 2)
Y = np.column_stack([X[:, 0], X[:, 1]])
@@ -1343,14 +1542,16 @@ def manually_create_model(equations, feature_names=None):
Path(output_directory)
/ run_id
/ f"hall_of_fame_output{i+1}.csv.bak"
- )
+ ),
+ index=False,
)
else:
model.nout_ = 1
model.selection_mask_ = None
model.feature_names_in_ = np.array(feature_names, dtype=object)
equations["complexity loss equation".split(" ")].to_csv(
- str(Path(output_directory) / run_id / "hall_of_fame.csv.bak")
+ str(Path(output_directory) / run_id / "hall_of_fame.csv.bak"),
+ index=False,
)
model.refresh()
@@ -1437,6 +1638,142 @@ def test_feature_selection_handler(self):
class TestMiscellaneous(unittest.TestCase):
"""Test miscellaneous functions."""
+ def test_search_releases_the_gil(self):
+ """Other Python threads must keep running while Julia searches."""
+ gaps = []
+ stop = threading.Event()
+
+ def sample_gaps():
+ previous = time.monotonic()
+ while not stop.is_set():
+ time.sleep(0.01)
+ now = time.monotonic()
+ gaps.append(now - previous)
+ previous = now
+
+ rstate = np.random.RandomState(0)
+ X = rstate.randn(100, 2)
+ y = X[:, 0] * X[:, 1]
+ model = PySRRegressor(
+ niterations=40,
+ populations=8,
+ verbosity=0,
+ progress=False,
+ temp_equation_file=True,
+ )
+
+ watcher = threading.Thread(target=sample_gaps, daemon=True)
+ watcher.start()
+ started = time.monotonic()
+ try:
+ model.fit(X, y)
+ finally:
+ elapsed = time.monotonic() - started
+ stop.set()
+ watcher.join(timeout=5)
+
+ # Holding the GIL across the search would leave one gap spanning it.
+ self.assertLess(max(gaps), elapsed / 4)
+
+ def test_loss_function_can_call_back_into_python(self):
+ """A Julia loss_function shim can drive a Python objective mid-search."""
+ calls = {"n": 0}
+
+ def python_objective(tree, dataset, options):
+ calls["n"] += 1
+ prediction, completed = jl.SymbolicRegression.eval_tree_array(
+ tree, dataset.X, options
+ )
+ if not completed:
+ return float("inf")
+ prediction = np.asarray(prediction)
+ target = np.asarray(dataset.y)
+ return float(np.mean((prediction - target) ** 2))
+
+ jl.pysr_test_python_objective = python_objective
+ jl.seval("""
+ using PythonCall
+
+ function pysr_test_python_objective_shim(
+ tree, dataset::Dataset{T,L}, options
+ )::L where {T,L}
+ PythonCall.GIL.@lock begin
+ return pyconvert(
+ L, pysr_test_python_objective(tree, dataset, options)
+ )
+ end
+ end
+ """)
+
+ X = np.arange(8.0).reshape(-1, 1)
+ y = 2.0 * X[:, 0]
+ model = PySRRegressor(
+ binary_operators=["+"],
+ unary_operators=[],
+ loss_function="pysr_test_python_objective_shim",
+ niterations=1,
+ populations=1,
+ population_size=20,
+ ncycles_per_iteration=1,
+ maxsize=7,
+ should_optimize_constants=False,
+ parallelism="serial",
+ deterministic=True,
+ random_state=0,
+ precision=64,
+ progress=False,
+ verbosity=0,
+ temp_equation_file=True,
+ )
+ model.fit(X, y)
+
+ self.assertGreater(calls["n"], 0)
+ self.assertTrue(np.all(np.isfinite(model.equations_["loss"])))
+
+ def test_unfitted_julia_streams_are_none(self):
+ model = PySRRegressor()
+ self.assertIsNone(model.julia_options_)
+ self.assertIsNone(model.julia_state_)
+
+ def test_inline_operator_must_be_a_function(self):
+ from pysr.sr import _maybe_create_inline_operators
+
+ with self.assertRaisesRegex(ValueError, "must evaluate to a Julia function"):
+ _maybe_create_inline_operators({1: ["identity(1)"]}, None, False)
+
+ def test_variable_names_must_match_feature_count(self):
+ # sklearn's own length check fires first during `fit`; exercise the
+ # backstop in `_check_assertions` directly.
+ from pysr.sr import _check_assertions
+
+ X = np.random.randn(10, 2)
+ with self.assertRaisesRegex(ValueError, "one name per feature"):
+ _check_assertions(X, True, ["a"], None, None, X[:, 0], None, None, False)
+
+ def test_equations_with_quoted_constants_are_read(self):
+ """Equations holding string constants arrive with doubled quotes."""
+ output_directory = tempfile.mkdtemp()
+ run_id = "quoted"
+ os.makedirs(Path(output_directory) / run_id)
+ equation = 'f = "a"; p = ["a", "b"]'
+ with open(
+ Path(output_directory) / run_id / "hall_of_fame.csv", "w", encoding="utf-8"
+ ) as f:
+ f.write("Complexity,Loss,Equation\n")
+ f.write('1,0.0,"""a"""\n')
+ f.write('3,0.0,"f = ""a""; p = [""a"", ""b""]"\n')
+
+ model = PySRRegressor()
+ model.nout_ = 1
+ model.output_directory_ = output_directory
+ model.run_id_ = run_id
+
+ equations = model._read_equation_file()[0]
+
+ self.assertEqual(list(equations["complexity"]), [1, 3])
+ self.assertEqual(list(equations["loss"]), [0.0, 0.0])
+ self.assertEqual(list(equations["equation"]), ['"a"', equation])
+
def test_pickle_inv_sympy_expression(self):
"""Test that sympy expressions with the inv operator can be pickled and unpickled correctly."""
expr_str = "inv(x0) + x1"
@@ -1518,6 +1855,104 @@ def test_pickle_with_temp_equation_file(self):
y_predictions2 = model2.predict(X)
np.testing.assert_array_almost_equal(y_predictions, y_predictions2)
+ def test_checkpoint_schema(self):
+ state = PySRRegressor().__getstate__()
+ self.assertEqual(state["_checkpoint_schema_version"], 3)
+
+ schema_two_state = state.copy()
+ schema_two_state.pop("type_spec")
+ schema_two_model = PySRRegressor.__new__(PySRRegressor)
+ schema_two_model.__setstate__(schema_two_state)
+ self.assertIsNone(schema_two_model.type_spec)
+
+ missing_schema_state = state.copy()
+ missing_schema_state.pop("_checkpoint_schema_version")
+ missing_schema_model = PySRRegressor.__new__(PySRRegressor)
+ with self.assertRaisesRegex(ValueError, "Unsupported PySR checkpoint schema"):
+ missing_schema_model.__setstate__(missing_schema_state)
+
+ for schema_version in (1, 2, 4):
+ incompatible_state = state.copy()
+ incompatible_state["_checkpoint_schema_version"] = schema_version
+
+ model = PySRRegressor.__new__(PySRRegressor)
+ with self.assertRaisesRegex(
+ ValueError, "Unsupported PySR checkpoint schema"
+ ):
+ model.__setstate__(incompatible_state)
+
+ def test_checkpoint_dump_failure_preserves_existing_file(self):
+ def fail_after_partial_write(_, checkpoint_file):
+ checkpoint_file.write(b"partial")
+ raise TypeError("cannot pickle")
+
+ with tempfile.TemporaryDirectory() as directory:
+ model = PySRRegressor()
+ model.output_directory_ = directory
+ model.run_id_ = "failed-checkpoint"
+ model.show_pickle_warnings_ = False
+ checkpoint = model.get_pkl_filename()
+ checkpoint.write_bytes(b"previous checkpoint")
+
+ with mock.patch("pysr.sr.pkl.dump", side_effect=fail_after_partial_write):
+ model._checkpoint()
+
+ self.assertEqual(checkpoint.read_bytes(), b"previous checkpoint")
+ self.assertFalse(model.show_pickle_warnings_)
+ self.assertEqual(list(checkpoint.parent.glob("*.tmp")), [])
+
+ def test_checkpoint_dump_failure_leaves_no_file(self):
+ def fail_after_partial_write(_, checkpoint_file):
+ checkpoint_file.write(b"partial")
+ raise TypeError("cannot pickle")
+
+ with tempfile.TemporaryDirectory() as directory:
+ model = PySRRegressor()
+ model.output_directory_ = directory
+ model.run_id_ = "failed-checkpoint"
+ checkpoint = model.get_pkl_filename()
+
+ with mock.patch("pysr.sr.pkl.dump", side_effect=fail_after_partial_write):
+ model._checkpoint()
+
+ self.assertFalse(checkpoint.exists())
+ self.assertTrue(model.show_pickle_warnings_)
+ self.assertEqual(list(checkpoint.parent.glob("*.tmp")), [])
+
+ def test_checkpoint_cleanup_failure_restores_warnings(self):
+ with tempfile.TemporaryDirectory() as directory:
+ model = PySRRegressor()
+ model.output_directory_ = directory
+ model.run_id_ = "failed-cleanup"
+ model.show_pickle_warnings_ = True
+ checkpoint = model.get_pkl_filename()
+
+ with mock.patch(
+ "pathlib.Path.unlink", side_effect=OSError("cannot unlink")
+ ):
+ model._checkpoint()
+
+ self.assertTrue(checkpoint.exists())
+ self.assertTrue(model.show_pickle_warnings_)
+
+ def test_checkpoint_propagates_publication_failure(self):
+ with tempfile.TemporaryDirectory() as directory:
+ model = PySRRegressor()
+ model.output_directory_ = directory
+ model.run_id_ = "unpublishable-checkpoint"
+ model.show_pickle_warnings_ = True
+ checkpoint = model.get_pkl_filename()
+
+ with mock.patch(
+ "pysr.sr.os.replace", side_effect=OSError("read-only directory")
+ ):
+ with self.assertRaises(OSError):
+ model._checkpoint()
+
+ self.assertFalse(checkpoint.exists())
+ self.assertEqual(list(checkpoint.parent.glob("*.tmp")), [])
+ self.assertTrue(model.show_pickle_warnings_)
+
def test_scikit_learn_compatibility(self):
"""Test PySRRegressor compatibility with scikit-learn."""
model = PySRRegressor(
@@ -1578,7 +2013,7 @@ def test_invalid_batch_size_corrects_and_warns(self):
self.assertTrue(any("batch_size" in str(w.message) for w in caught))
- def test_progress_disabled_when_stdout_lacks_buffer(self):
+ def test_progress_silently_disabled_when_stdout_lacks_buffer(self):
fake_stdout = type(
"FakeStdout", (), {"write": lambda self, *_args, **_kwargs: None}
)()
@@ -1596,7 +2031,286 @@ def test_progress_disabled_when_stdout_lacks_buffer(self):
y = X[:, 0]
model.fit(X, y)
- self.assertTrue(any("progress bar" in str(w.message) for w in caught))
+ self.assertFalse(any("progress bar" in str(w.message) for w in caught))
+
+ def test_builtin_mutation_and_plugin_configs(self):
+ from pysr import (
+ AdaptiveMutationWeightsPlugin,
+ BacksolveMutation,
+ ConstantMutation,
+ MutationBurstPlugin,
+ OperatorMutation,
+ SymbolicRegression,
+ )
+
+ mutation_name = jl.seval("m -> string(nameof(typeof(m)))")
+ for mutation, expected_name in (
+ (OperatorMutation(), "OperatorMutation"),
+ (BacksolveMutation(), "BacksolveMutation"),
+ ):
+ julia_mutation = mutation.julia_mutation()
+ self.assertEqual(str(mutation_name(julia_mutation)), expected_name)
+
+ backsolve_mutation = BacksolveMutation(
+ max_library_size=321,
+ max_terms=4,
+ min_improvement=0.02,
+ node_attempts=5,
+ ).julia_mutation()
+ self.assertEqual(backsolve_mutation.max_library_size, 321)
+ self.assertEqual(backsolve_mutation.max_terms, 4)
+ self.assertEqual(backsolve_mutation.min_improvement, 0.02)
+ self.assertEqual(backsolve_mutation.node_attempts, 5)
+ constant_mutation = ConstantMutation().julia_mutation()
+ backend_default = SymbolicRegression.ConstantMutation()
+ self.assertEqual(
+ constant_mutation.perturbation_factor,
+ backend_default.perturbation_factor,
+ )
+ self.assertEqual(
+ constant_mutation.probability_negate,
+ backend_default.probability_negate,
+ )
+
+ plugin_name = jl.seval("p -> string(nameof(typeof(p)))")
+ for plugin, expected_name in (
+ (
+ AdaptiveMutationWeightsPlugin(reward="loss"),
+ "AdaptiveMutationWeightsPlugin",
+ ),
+ (MutationBurstPlugin(retry_attempts=2), "MutationBurstPlugin"),
+ ):
+ self.assertEqual(str(plugin_name(plugin.julia_plugin())), expected_name)
+
+ def test_default_operators_and_plugins_match_backend(self):
+ from pysr import SymbolicRegression
+
+ model = PySRRegressor()
+ backend_options = SymbolicRegression.Options()
+
+ backend_binary_operators = jl.seval(
+ "options -> string.(options.operators.ops[2])"
+ )(backend_options)
+ self.assertListEqual(
+ model._validate_and_modify_params().operators[2],
+ [str(operator) for operator in backend_binary_operators],
+ )
+
+ plugin_name = jl.seval("p -> string(nameof(typeof(p)))")
+ plugins = {
+ str(plugin_name(plugin)): plugin for plugin in backend_options.plugins
+ }
+
+ parsimony = plugins["AdaptiveParsimonyPlugin"]
+ self.assertEqual(model.use_frequency, parsimony.mutation_acceptance)
+ self.assertEqual(model.use_frequency_in_tournament, parsimony.tournament)
+
+ def test_mutation_and_plugin_configuration(self):
+ from pysr import (
+ AbstractPlugin,
+ AdaptiveParsimonyPlugin,
+ BacksolveMutation,
+ ConstantMutation,
+ RandomizeMutation,
+ SimulatedAnnealingPlugin,
+ SymbolicRegression,
+ )
+
+ jl.seval("""
+ if !isdefined(Main, :PySRPluginMutationTest)
+ @eval module PySRPluginMutationTest
+ using SymbolicRegression
+ struct PluginMutation <: SymbolicRegression.AbstractMutation end
+ struct Plugin <: SymbolicRegression.AbstractPlugin end
+ SymbolicRegression.plugin_mutations(::Plugin) =
+ (PluginMutation() => 0.25,)
+ end
+ end
+ """)
+
+ class PluginConfig(AbstractPlugin):
+ def julia_plugin(self):
+ return jl.seval("PySRPluginMutationTest.Plugin()")
+
+ X = np.arange(12, dtype=np.float32).reshape(6, 2)
+ y = X[:, 0]
+ model = PySRRegressor(
+ binary_operators=["+"],
+ unary_operators=[],
+ niterations=0,
+ populations=1,
+ population_size=4,
+ tournament_selection_n=2,
+ progress=False,
+ temp_equation_file=True,
+ perturbation_factor=0.7,
+ probability_negate_constant=0.8,
+ mutations={
+ BacksolveMutation(max_library_size=123): 0.02,
+ ConstantMutation(perturbation_factor=0.2): 0.5,
+ RandomizeMutation(): 0.4,
+ },
+ plugins=[
+ SimulatedAnnealingPlugin(alpha=0.4),
+ AdaptiveParsimonyPlugin(tournament=False, mutation_acceptance=False),
+ PluginConfig(),
+ ],
+ )
+ model.fit(X, y)
+
+ options = model.julia_options_
+ mutation_name = jl.seval("p -> string(nameof(typeof(first(p))))")
+ mutation_by_name = {
+ str(mutation_name(pair)): pair for pair in options.mutations
+ }
+ self.assertEqual(options.crossover_probability, 0.2)
+ expected_weights = {
+ str(mutation_name(pair)): float(jl.last(pair))
+ for pair in SymbolicRegression.Options().mutations
+ }
+ expected_weights.update(
+ ConstantMutation=0.5,
+ RandomizeMutation=0.4,
+ BacksolveMutation=0.02,
+ PluginMutation=0.25,
+ )
+ self.assertSetEqual(set(mutation_by_name), set(expected_weights))
+ for name, weight in expected_weights.items():
+ self.assertEqual(float(jl.last(mutation_by_name[name])), weight)
+ constant_mutation = jl.first(mutation_by_name["ConstantMutation"])
+ self.assertEqual(constant_mutation.perturbation_factor, 0.2)
+ self.assertEqual(
+ constant_mutation.probability_negate,
+ SymbolicRegression.ConstantMutation().probability_negate,
+ )
+ self.assertEqual(
+ jl.first(mutation_by_name["BacksolveMutation"]).max_library_size,
+ 123,
+ )
+
+ plugin_name = jl.seval("p -> string(nameof(typeof(p)))")
+ plugins = {str(plugin_name(plugin)): plugin for plugin in options.plugins}
+ self.assertSetEqual(
+ set(plugins),
+ {
+ "SimulatedAnnealingPlugin",
+ "AdaptiveParsimonyPlugin",
+ "AdaptiveMutationWeightsPlugin",
+ "Plugin",
+ },
+ )
+ self.assertEqual(plugins["SimulatedAnnealingPlugin"].alpha, 0.4)
+ self.assertFalse(plugins["AdaptiveParsimonyPlugin"].tournament)
+ self.assertFalse(plugins["AdaptiveParsimonyPlugin"].mutation_acceptance)
+
+ def test_default_mutation_and_plugin_configuration(self):
+ from pysr import (
+ BacksolveMutation,
+ MutationBurstPlugin,
+ OperatorMutation,
+ SimulatedAnnealingPlugin,
+ )
+
+ X = np.arange(12, dtype=np.float32).reshape(6, 2)
+ y = X[:, 0]
+ model = PySRRegressor(
+ binary_operators=["+"],
+ unary_operators=[],
+ niterations=0,
+ populations=1,
+ population_size=4,
+ tournament_selection_n=2,
+ progress=False,
+ temp_equation_file=True,
+ default_mutations={OperatorMutation(): 0.7},
+ mutations={BacksolveMutation(max_library_size=123): 0.2},
+ default_plugins=[MutationBurstPlugin(retry_attempts=2)],
+ plugins=[SimulatedAnnealingPlugin(alpha=0.4)],
+ )
+ model.fit(X, y)
+
+ mutation_name = jl.seval("p -> string(nameof(typeof(first(p))))")
+ weights = {
+ str(mutation_name(pair)): float(jl.last(pair))
+ for pair in model.julia_options_.mutations
+ }
+ self.assertDictEqual(
+ weights,
+ {"OperatorMutation": 0.7, "BacksolveMutation": 0.2},
+ )
+
+ plugin_name = jl.seval("p -> string(nameof(typeof(p)))")
+ plugins = {
+ str(plugin_name(plugin)): plugin for plugin in model.julia_options_.plugins
+ }
+ self.assertSetEqual(
+ set(plugins), {"MutationBurstPlugin", "SimulatedAnnealingPlugin"}
+ )
+ self.assertEqual(plugins["MutationBurstPlugin"].retry_attempts, 2)
+ self.assertEqual(plugins["SimulatedAnnealingPlugin"].alpha, 0.4)
+
+ def test_legacy_mutation_weight_override(self):
+ from pysr import SymbolicRegression
+
+ X = np.arange(12, dtype=np.float32).reshape(6, 2)
+ y = X[:, 0]
+ model = PySRRegressor(
+ binary_operators=["+"],
+ unary_operators=[],
+ niterations=0,
+ populations=1,
+ population_size=4,
+ tournament_selection_n=2,
+ progress=False,
+ temp_equation_file=True,
+ perturbation_factor=0.7,
+ probability_negate_constant=0.8,
+ weight_randomize=0.3,
+ weight_backsolve=0.2,
+ )
+ model.fit(X, y)
+
+ mutation_name = jl.seval("p -> string(nameof(typeof(first(p))))")
+ weights = {
+ str(mutation_name(pair)): float(jl.last(pair))
+ for pair in model.julia_options_.mutations
+ }
+ self.assertEqual(weights["RandomizeMutation"], 0.3)
+ self.assertEqual(weights["BacksolveMutation"], 0.2)
+ backend_weights = {
+ str(mutation_name(pair)): float(jl.last(pair))
+ for pair in SymbolicRegression.default_mutations()
+ }
+ self.assertEqual(weights["AddNodeMutation"], backend_weights["AddNodeMutation"])
+ constant_mutation = next(
+ jl.first(pair)
+ for pair in model.julia_options_.mutations
+ if str(mutation_name(pair)) == "ConstantMutation"
+ )
+ self.assertEqual(constant_mutation.perturbation_factor, 0.7)
+ self.assertEqual(constant_mutation.probability_negate, 0.8)
+
+ def test_mutation_interfaces_are_mutually_exclusive(self):
+ from sklearn.base import clone
+
+ from pysr import OperatorMutation
+
+ model = PySRRegressor()
+ cloned_model = clone(model)
+ legacy_parameters = filter(
+ lambda parameter: parameter.startswith("weight_"), model.get_params()
+ )
+ for parameter in legacy_parameters:
+ self.assertIsNone(getattr(model, parameter))
+ self.assertIsNone(getattr(cloned_model, parameter))
+
+ for new_interface in (
+ {"default_mutations": {OperatorMutation(): 0.7}},
+ {"mutations": {OperatorMutation(): 0.7}},
+ ):
+ conflicting_model = PySRRegressor(weight_add_node=2.47, **new_interface)
+ with self.assertRaisesRegex(ValueError, "Cannot combine legacy"):
+ conflicting_model._validate_and_modify_params()
def test_param_groupings(self):
"""Test that param_groupings are complete"""
@@ -1621,54 +2335,18 @@ def test_param_groupings(self):
# Check the sets are equal:
self.assertSetEqual(set(params), set(regressor_params))
- def test_parametric_deprecation_warning(self):
- """Test that the helpful warning message is displayed."""
- pattern = re.compile(
- r"ParametricExpressionSpec is deprecated.*TemplateExpressionSpec.*"
- r"max_parameters=2.*"
- r"variable_names=\[\"alpha\", \"beta\"\].*"
- r"expressions=\[\"f\"\].*"
- r"variable_names=\[\"alpha\", \"beta\", \"category\"\].*"
- r"parameters=\{\s*\"p1\": n_categories,\s*\"p2\": n_categories\s*\}.*"
- r"combine=\"f\(alpha, beta, p1\[category\], p2\[category\]\)\"",
- flags=re.S,
- )
-
- with self.assertWarnsRegex(FutureWarning, pattern):
- parametric_expression_deprecation_warning(
- max_parameters=2,
- variable_names=["alpha", "beta"],
- )
-
def test_load_all_packages(self):
"""Test we can load all packages at once."""
load_all_packages()
self.assertTrue(jl.seval("ClusterManagers isa Module"))
+ self.assertTrue(jl.seval("SlurmClusterManager isa Module"))
- def test_get_batch_size(self):
- """Test the _get_batch_size function."""
- from pysr.sr import _get_batch_size
-
- # Test None (auto) mode with different dataset sizes
- self.assertEqual(_get_batch_size(500, None), 500)
- self.assertEqual(_get_batch_size(999, None), 999)
- self.assertEqual(_get_batch_size(1000, None), 1000)
- self.assertEqual(_get_batch_size(1001, None), 128)
- self.assertEqual(_get_batch_size(1500, None), 128)
- self.assertEqual(_get_batch_size(4999, None), 128)
- self.assertEqual(_get_batch_size(5000, None), 256)
- self.assertEqual(_get_batch_size(10000, None), 256)
- self.assertEqual(_get_batch_size(49999, None), 256)
- self.assertEqual(_get_batch_size(50000, None), 512)
- self.assertEqual(_get_batch_size(100000, None), 512)
-
- # Test explicit batch_size
- self.assertEqual(_get_batch_size(1000, 64), 64)
- self.assertEqual(_get_batch_size(1000, 2000), 1000) # Capped at dataset size
- self.assertEqual(_get_batch_size(50, 100), 50) # Capped at dataset size
+ with mock.patch.dict(os.environ, {"SLURM_JOB_ID": "1", "SLURM_NTASKS": "2"}):
+ with self.assertRaisesRegex(JuliaError, "allocation has 2 tasks"):
+ _load_cluster_manager("slurm")(1)
def test_batching_auto(self):
- """Test that batching='auto' works correctly."""
+ """Test that batching configuration is passed to SymbolicRegression.jl."""
model = PySRRegressor()
self.assertEqual(model.batching, "auto")
@@ -1676,11 +2354,13 @@ def test_batching_auto(self):
y_small = np.random.randn(100)
model = PySRRegressor(batching="auto", niterations=0)
model.fit(X_small, y_small)
+ self.assertEqual(model.julia_options_.batching, jl.Symbol("auto"))
+ self.assertIsNone(model.julia_options_.batch_size)
- X_large = np.random.randn(1001, 2)
- y_large = np.random.randn(1001)
- model2 = PySRRegressor(batching="auto", niterations=0)
- model2.fit(X_large, y_large)
+ explicit = PySRRegressor(batching=True, batch_size=64, niterations=0)
+ explicit.fit(X_small, y_small)
+ self.assertTrue(explicit.julia_options_.batching)
+ self.assertEqual(explicit.julia_options_.batch_size, 64)
def test_batch_size_negative_warning(self):
"""Test that batch_size < 1 gives a warning for integers only."""
@@ -1752,17 +2432,6 @@ def test_from_file_requires_operator_configuration(self):
self.assertIn("must provide either `operators`", str(cm.exception))
- def test_size_warning(self):
- """Ensure that a warning is given for a large input size."""
- model = PySRRegressor()
- X = np.random.randn(50001, 2)
- y = np.random.randn(50001)
- with warnings.catch_warnings():
- warnings.simplefilter("error")
- with self.assertRaises(Exception) as context:
- model.fit(X, y)
- self.assertIn("more than 50,000", str(context.exception))
-
def test_deterministic_warnings(self):
"""Ensure that warnings are given for determinism"""
model = PySRRegressor(random_state=0)
@@ -2150,6 +2819,7 @@ def test_unit_checks(self):
y,
X_units,
y_units,
+ supports_sympy=False,
)
invalid_units = [
(np.ones((10, 2)), np.ones(10), ["m/s", "s", "s^2"], None),
@@ -2165,6 +2835,7 @@ def test_unit_checks(self):
y,
X_units,
y_units,
+ supports_sympy=False,
)
def test_unit_propagation(self):
@@ -2246,22 +2917,20 @@ def test_process_constraints_swaps_multiplication_constraints(self):
class TestTemplateExpressionSpec(unittest.TestCase):
- def test_num_features_symbol_keys(self):
- # ponytail: one check — dict keys must reach Julia as Symbols
- spec = TemplateExpressionSpec(
- ["f", "g"],
- "combine(fs, vars) = fs.f(vars[1], vars[2]) + fs.g(vars[3])",
- {"f": 2, "g": 1},
- )
- options = spec.julia_expression_options()
- names = jl.seval("x -> propertynames(x.structure.num_features)")(options)
- self.assertEqual(names, (jl.Symbol("f"), jl.Symbol("g")))
-
def _check_macro_str(self, spec, expected_str):
self.assertEqual(
spec._template_macro_str().strip(), dedent(expected_str).strip()
)
+ def test_is_dataclass(self):
+ spec = TemplateExpressionSpec("f(x)", ["f"], ["x"])
+
+ self.assertTrue(is_dataclass(spec))
+ self.assertEqual(
+ [field.name for field in fields(spec)],
+ ["combine", "expressions", "variable_names", "parameters"],
+ )
+
def test_single_expression_no_params_single_variable(self):
spec = TemplateExpressionSpec(
combine="f(x)", expressions=["f"], variable_names=["x"]
@@ -2382,6 +3051,7 @@ def runtests(just_tests=False):
TestLaTeXTable,
TestDimensionalConstraints,
TestGuesses,
+ TestTypeSpecs,
]
if just_tests:
return test_cases
diff --git a/pysr/test/test_nb.ipynb b/pysr/test/test_nb.ipynb
index 0a2050a2a..f2b93382c 100644
--- a/pysr/test/test_nb.ipynb
+++ b/pysr/test/test_nb.ipynb
@@ -4,19 +4,12 @@
"cell_type": "code",
"execution_count": 1,
"metadata": {},
- "outputs": [
- {
- "name": "stdout",
- "output_type": "stream",
- "text": [
- "Detected Jupyter notebook. Loading juliacall extension. Set `PYSR_AUTOLOAD_EXTENSIONS=no` to disable.\n"
- ]
- }
- ],
+ "outputs": [],
"source": [
"# NBVAL_IGNORE_OUTPUT\n",
"import numpy as np\n",
- "from pysr import PySRRegressor, jl"
+ "from pysr import PySRRegressor, jl\n",
+ "%load_ext juliacall"
]
},
{
@@ -35,7 +28,7 @@
"source": [
"%%julia\n",
"\n",
- "# Automatically activates Julia magic\n",
+ "# Julia magic comes from the juliacall extension loaded above\n",
"\n",
"x = 1\n",
"println(x + 2)"
diff --git a/pysr/test/test_slurm.py b/pysr/test/test_slurm.py
index bda248c91..eae0679a8 100644
--- a/pysr/test/test_slurm.py
+++ b/pysr/test/test_slurm.py
@@ -154,9 +154,10 @@ def test_pysr_slurm_cluster_manager(self):
"#SBATCH --job-name=pysr-slurm-test",
"#SBATCH --partition=normal",
"#SBATCH --nodes=2",
- "#SBATCH --ntasks=4",
+ "#SBATCH --ntasks-per-node=2",
"#SBATCH --time=40:00",
"set -euo pipefail",
+ "export JULIA_DEBUG=SlurmClusterManager",
"python3 - <<'PY'",
"import numpy as np",
"from pysr import PySRRegressor",
@@ -164,11 +165,11 @@ def test_pysr_slurm_cluster_manager(self):
"y = X[:, 0] + 1.0",
"model = PySRRegressor(",
" niterations=3,",
- " populations=3,",
+ " populations=4,",
" progress=False,",
" temp_equation_file=True,",
" parallelism='multiprocessing',",
- " procs=2,",
+ " procs=4,",
" cluster_manager='slurm',",
" verbosity=0,",
")",
@@ -275,6 +276,15 @@ def test_pysr_slurm_cluster_manager(self):
env=self.compose_env,
)
self.assertEqual(output.returncode, 0, msg=output.stdout)
+ worker_hosts = re.findall(
+ r"Worker \d+ ready on host ([^,]+), port \d+", output.stdout
+ )
+ self.assertEqual(len(worker_hosts), 4, msg=output.stdout)
+ self.assertEqual(
+ sorted(worker_hosts.count(host) for host in set(worker_hosts)),
+ [2, 2],
+ msg=output.stdout,
+ )
self.assertIn("PYSR_SLURM_OK", output.stdout)
diff --git a/pysr/test/test_startup.py b/pysr/test/test_startup.py
index 8bd81cd93..9a0686d42 100644
--- a/pysr/test/test_startup.py
+++ b/pysr/test/test_startup.py
@@ -131,6 +131,86 @@ def test_bad_startup_options(self):
)
self.assertIn(warning_test["msg"], result.stderr.decode())
+ def test_autoload_extension_env_precedence(self):
+ autoload_key = "PYTHON_JULIACALL_AUTOLOAD_IPYTHON_EXTENSION"
+ deprecated_key = "PYSR_AUTOLOAD_EXTENSIONS"
+ cases = [
+ # (extra env, expected value after `import pysr`)
+ ({}, "no"),
+ ({autoload_key: "yes"}, "yes"),
+ ({deprecated_key: "yes"}, "yes"),
+ # An explicitly set juliacall var beats the deprecated alias:
+ ({autoload_key: "no", deprecated_key: "yes"}, "no"),
+ ]
+ for extra_env, expected in cases:
+ env = {
+ k: v
+ for k, v in os.environ.items()
+ if k not in (autoload_key, deprecated_key)
+ }
+ env.update(extra_env)
+ result = subprocess.run(
+ [
+ sys.executable,
+ "-c",
+ f'import pysr, os; print("autoload=" + os.environ["{autoload_key}"])',
+ ],
+ stdout=subprocess.PIPE,
+ stderr=subprocess.PIPE,
+ env=env,
+ )
+ self.assertEqual(result.returncode, 0, result.stderr.decode())
+ self.assertIn(f"autoload={expected}", result.stdout.decode())
+ deprecation_expected = deprecated_key in extra_env
+ self.assertEqual(
+ "PYSR_AUTOLOAD_EXTENSIONS is deprecated" in result.stderr.decode(),
+ deprecation_expected,
+ )
+
+ def test_juliacall_0935_distributed_worker_uses_current_python(self):
+ import_orders = (
+ "import pysr\nfrom pysr import jl",
+ "import juliacall\nimport pysr\nfrom pysr import jl",
+ )
+ worker_code = r"""
+ using Distributed
+ worker = only(addprocs(1, exeflags="--threads=1"))
+ try
+ fetch(Distributed.remotecall_eval(Main, worker, :(using PythonCall)))
+ @fetchfrom worker (
+ pyconvert(String, pyimport("sys").executable),
+ pyconvert(Int, pybuiltins.sum([1, 2, 3])),
+ )
+ finally
+ rmprocs(worker)
+ end
+ """
+ env = os.environ.copy()
+ env.pop("JULIA_PYTHONCALL_EXE", None)
+ env.pop("JULIA_PYTHONCALL_EXECUTABLE", None)
+ env["JULIA_CONDAPKG_OFFLINE"] = "yes"
+ env["PYTHON_JULIACALL_THREADS"] = "1"
+ env["JULIA_NUM_THREADS"] = "1"
+
+ for import_order in import_orders:
+ with self.subTest(import_order=import_order):
+ code = (
+ "import sys\n"
+ f"{import_order}\n"
+ f"worker_python, total = jl.seval({worker_code!r})\n"
+ "assert worker_python == sys.executable, "
+ "(worker_python, sys.executable)\n"
+ "assert total == 6, total\n"
+ )
+ result = subprocess.run(
+ [sys.executable, "-c", code],
+ capture_output=True,
+ text=True,
+ env=env,
+ timeout=300,
+ )
+ self.assertEqual(result.returncode, 0, result.stdout + result.stderr)
+
def test_notebook(self):
if platform.system() == "Windows":
self.skipTest("Notebook test incompatible with Windows")
diff --git a/pysr/test/test_torch.py b/pysr/test/test_torch.py
index e66e7d960..0d0933d0f 100644
--- a/pysr/test/test_torch.py
+++ b/pysr/test/test_torch.py
@@ -51,7 +51,7 @@ def test_pipeline_pandas(self):
for fname in ["hall_of_fame.csv.bak", "hall_of_fame.csv"]:
equations["Complexity Loss Equation".split(" ")].to_csv(
- Path(model.output_directory_) / model.run_id_ / fname
+ Path(model.output_directory_) / model.run_id_ / fname, index=False
)
model.refresh(run_directory=str(Path(model.output_directory_) / model.run_id_))
@@ -85,7 +85,7 @@ def test_pipeline(self):
for fname in ["hall_of_fame.csv.bak", "hall_of_fame.csv"]:
equations["Complexity Loss Equation".split(" ")].to_csv(
- Path(model.output_directory_) / model.run_id_ / fname
+ Path(model.output_directory_) / model.run_id_ / fname, index=False
)
model.refresh(run_directory=str(Path(model.output_directory_) / model.run_id_))
@@ -138,7 +138,7 @@ def test_custom_operator(self):
for fname in ["hall_of_fame.csv.bak", "hall_of_fame.csv"]:
equations["Complexity Loss Equation".split(" ")].to_csv(
- Path(model.output_directory_) / model.run_id_ / fname
+ Path(model.output_directory_) / model.run_id_ / fname, index=False
)
MyCustomOperator = sympy.Function("mycustomoperator")
@@ -194,7 +194,7 @@ def test_issue_656(self):
def test_constant_arguments(self):
# Test that functions with constant arguments work correctly
- # Regression test for https://github.com/MilesCranmer/PySR/issues/656
+ # Regression test for https://github.com/astroautomata/PySR/issues/656
test_cases = [
(pysr.export_sympy.pysr2sympy("sqrt(2)"), np.sqrt(2)),
(sympy.exp(2), np.exp(2)),
@@ -246,8 +246,7 @@ def cos_approx(x):
torch_output = torch_module(self.torch.tensor(X.values)).detach().numpy()
- np.testing.assert_almost_equal(y.values, np_output, decimal=3)
- np.testing.assert_almost_equal(y.values, torch_output, decimal=3)
+ np.testing.assert_almost_equal(np_output, torch_output, decimal=3)
def runtests(just_tests=False):
diff --git a/pysr/test/test_type_specs.py b/pysr/test/test_type_specs.py
new file mode 100644
index 000000000..e3741a0b8
--- /dev/null
+++ b/pysr/test/test_type_specs.py
@@ -0,0 +1,1797 @@
+import json
+import pickle
+import subprocess
+import sys
+import tempfile
+import unittest
+import uuid
+import warnings
+from pathlib import Path
+from unittest.mock import patch
+
+import numpy as np
+import pandas as pd
+from juliacall import JuliaError # type: ignore
+
+from pysr import PySRRegressor, TemplateExpressionSpec, TypeSpec, jl
+from pysr.expression_specs import AbstractExpressionSpec, ExpressionSpec
+from pysr.type_specs import (
+ CallableJuliaExpression,
+ compile_type_spec,
+ compile_type_spec_runtime,
+ load_type_spec_runtime,
+ object_array_1d,
+ object_array_2d,
+ prepare_type_spec_fit_data,
+ prepare_type_spec_prediction_data,
+ type_spec_to_julia_array,
+ type_spec_to_python_array,
+ validate_type_spec_configuration,
+ validate_type_spec_options,
+ validate_type_spec_runtime,
+)
+
+
+def string_spec(**overrides):
+ name = overrides.pop("name", "StringValue")
+ return TypeSpec(
+ name,
+ fields=overrides.pop("fields", {"data": "String"}),
+ sample=overrides.pop("sample", f'rng -> {name}(rand(rng, ("a", "b")))'),
+ mutate=overrides.pop(
+ "mutate",
+ f'(rng, value, temperature) -> {name}(rand(rng, ("a", "b")))',
+ ),
+ **overrides,
+ )
+
+
+def vector_spec(**overrides):
+ name = overrides.pop("name", "VectorValue")
+ return TypeSpec(
+ name,
+ fields=overrides.pop("fields", {"data": "Vector{Float64}"}),
+ sample=overrides.pop("sample", f"rng -> {name}([3.0, 4.0])"),
+ scalar_constants=overrides.pop("scalar_constants", "value -> value.data"),
+ with_scalar_constants=overrides.pop(
+ "with_scalar_constants",
+ f"(value, scalar_constants) -> {name}(collect(scalar_constants))",
+ ),
+ **overrides,
+ )
+
+
+def tiny_model(spec, *, parallelism="serial", **overrides):
+ type_name = spec.name
+ parameters = {
+ "type_spec": spec,
+ "operators": {1: [f"identity_{type_name}(x::{type_name}) = x"]},
+ "elementwise_loss": (
+ f"loss_{type_name}(x::{type_name}, y::{type_name})::Float64 = "
+ "x == y ? 0.0 : 1.0"
+ ),
+ "niterations": 1,
+ "ncycles_per_iteration": 2,
+ "populations": 1,
+ "population_size": 8,
+ "tournament_selection_n": 3,
+ "maxsize": 7,
+ "parallelism": parallelism,
+ "deterministic": parallelism == "serial",
+ "random_state": 0 if parallelism == "serial" else None,
+ "progress": False,
+ "verbosity": 0,
+ "temp_equation_file": True,
+ "should_optimize_constants": False,
+ }
+ parameters.update(overrides)
+ return PySRRegressor(**parameters)
+
+
+def identity_template():
+ return TemplateExpressionSpec(
+ combine="f(x)",
+ expressions=["f"],
+ variable_names=["x"],
+ )
+
+
+def string_data(*, constant: bool = False):
+ X = np.array([["a"], ["b"], ["a"], ["b"]], dtype=object)
+ y = np.full(len(X), "a", dtype=object) if constant else X[:, 0].copy()
+ return X, y
+
+
+class TestTypeSpecs(unittest.TestCase):
+ def test_init_invalid_value_hook(self):
+ spec = string_spec(
+ init_invalid='() -> StringValue("")',
+ is_valid="value -> !isempty(value.data)",
+ )
+ runtime = load_type_spec_runtime(compile_type_spec(spec))
+
+ self.assertTrue(
+ jl.seval(
+ "T -> begin v = SymbolicRegression.InterfaceDynamicExpressionsModule.DE.invalid_value(T); v isa T && isempty(v.data) end"
+ )(runtime.value_type)
+ )
+ X, y = string_data()
+ model = tiny_model(spec)
+ model.fit(X, y)
+ self.assertEqual(model.predict(X).tolist(), y.tolist())
+
+ def test_init_invalid_hook_rejects_wrong_type_or_valid_value(self):
+ for source, message in (
+ ('() -> ""', "init_invalid.*must return `StringValue`"),
+ (
+ '() -> StringValue("valid")',
+ "init_invalid.*must return an invalid value",
+ ),
+ ):
+ with self.subTest(source=source):
+ runtime = load_type_spec_runtime(
+ compile_type_spec(
+ string_spec(
+ init_invalid=source,
+ is_valid="value -> !isempty(value.data)",
+ )
+ )
+ )
+ with self.assertRaisesRegex(ValueError, message):
+ validate_type_spec_runtime(runtime)
+
+ def test_serial_string_type_spec(self):
+ X, y = string_data()
+ model = tiny_model(string_spec())
+ model.fit(X, y)
+
+ np.testing.assert_array_equal(model.predict(X), y)
+ self.assertTrue(
+ model._type_spec_runtime_definition_.module_name.startswith("_PySRConfig_")
+ )
+ self.assertIn("lambda_format", model.equations_.columns)
+
+ def test_named_type_is_imported_into_main(self):
+ suffix = uuid.uuid4().hex
+ type_name = f"MainVisibleValue_{suffix}"
+ runtime = load_type_spec_runtime(compile_type_spec(string_spec(name=type_name)))
+
+ self.assertTrue(bool(jl.isdefined(jl.Main, jl.Symbol(type_name))))
+ self.assertTrue(
+ bool(
+ jl.seval("(actual, expected) -> actual === expected")(
+ jl.getproperty(jl.Main, jl.Symbol(type_name)),
+ runtime.value_type,
+ )
+ )
+ )
+
+ def test_serial_vector_type_spec_with_scalar_constants(self):
+ spec = vector_spec()
+ X = np.empty((4, 1), dtype=object)
+ X[:, 0] = [
+ np.array([1.0, 2.0]),
+ np.array([3.0, 4.0]),
+ np.array([5.0, 6.0]),
+ np.array([7.0, 8.0]),
+ ]
+ model = tiny_model(spec)
+ model.fit(X, X[:, 0].copy())
+
+ self.assertEqual(
+ [list(value) for value in model.predict(X)],
+ [list(value) for value in X[:, 0]],
+ )
+ runtime = model._load_type_spec_runtime()
+ value = runtime.module._convert_value([1.0, 2.0])
+ packed, next_index = jl.seval("""
+ function (value)
+ DE = SymbolicRegression.InterfaceDynamicExpressionsModule.DE
+ buffer = fill(-99.0, 6)
+ next_index = DE.pack_scalar_constants!(buffer, 3, value)
+ return buffer, next_index
+ end
+ """)(value)
+ self.assertEqual(list(packed), [-99.0, -99.0, 1.0, 2.0, -99.0, -99.0])
+ self.assertEqual(int(next_index), 5)
+
+ def test_multithreaded_type_spec(self):
+ X, y = string_data()
+ model = tiny_model(string_spec(), parallelism="multithreading", niterations=10)
+ model.fit(X, y)
+ np.testing.assert_array_equal(model.predict(X), y)
+
+ def test_configuration_modules_isolate_colliding_operator_names(self):
+ X, y = string_data()
+ type_name = "IsolatedConfigurationValue"
+ operator_name = "colliding_identity"
+ first = tiny_model(
+ string_spec(name=type_name),
+ operators={1: [f"{operator_name}(x::{type_name}) = x"]},
+ )
+ second = tiny_model(
+ string_spec(name=type_name),
+ operators={1: [f'{operator_name}(x::{type_name}) = {type_name}("a")']},
+ )
+
+ first.fit(X, y)
+ second.fit(X, np.full(len(y), "a", dtype=object))
+
+ np.testing.assert_array_equal(first.predict(X), y)
+ np.testing.assert_array_equal(
+ second.predict(X), np.full(len(y), "a", dtype=object)
+ )
+ self.assertIsNot(
+ first._load_type_spec_runtime().configuration_module,
+ second._load_type_spec_runtime().configuration_module,
+ )
+
+ def test_preamble_struct_has_one_identity_in_hooks_and_operators(self):
+ type_name = "PreambleStructValue"
+ spec = TypeSpec(
+ type_name,
+ fields={"data": "PrivatePayload"},
+ preamble="struct PrivatePayload\n data::String\n end",
+ sample=f'rng -> {type_name}(PrivatePayload("a"))',
+ mutate="(rng, value, temperature) -> value",
+ )
+ definition = compile_type_spec_runtime(
+ spec,
+ {1: [f"x -> {type_name}(PrivatePayload(x.data.data))"]},
+ elementwise_loss="(prediction, target) -> 0.0",
+ loss_function=None,
+ loss_function_expression=None,
+ complexity_mapping=None,
+ early_stop_condition=None,
+ )
+ runtime = load_type_spec_runtime(definition)
+
+ same_object = jl.seval("(a, b) -> a === b")
+ self.assertTrue(
+ bool(
+ same_object(
+ jl.getproperty(runtime.module, jl.Symbol("PrivatePayload")),
+ jl.getproperty(
+ runtime.configuration_module, jl.Symbol("PrivatePayload")
+ ),
+ )
+ )
+ )
+ value = runtime.module._sample(runtime.module.Random.default_rng())
+ result = runtime.operator_functions[1][0](value)
+ self.assertTrue(bool(jl.isa(result.data, runtime.module.PrivatePayload)))
+
+ def test_definitions_support_fit_and_predict(self):
+ spec = TypeSpec(
+ "TransformedValue",
+ fields={"value": "Component"},
+ preamble="const Component = Float64",
+ definitions="component(value::TransformedValue) = value.value",
+ init="() -> TransformedValue(0.0)",
+ sample="rng -> TransformedValue(1.0)",
+ mutate="(rng, value, temperature) -> value",
+ )
+ X = object_array_2d([[0.0], [1.0], [2.0], [3.0]])
+ y = object_array_1d([2.0, 4.0, 6.0, 8.0])
+ model = tiny_model(
+ spec,
+ operators={
+ 1: [
+ """
+ transform(value::TransformedValue)::TransformedValue =
+ TransformedValue(2 * (component(value) + 1))
+ """
+ ]
+ },
+ elementwise_loss="""
+ transformed_loss(a::TransformedValue, b::TransformedValue)::Float64 =
+ abs2(component(a) - component(b))
+ """,
+ guesses=["transform(x0)"],
+ )
+ model.fit(X, y)
+ np.testing.assert_array_equal(model.predict(X), y)
+
+ def test_definitions_are_private_to_their_specification(self):
+ type_name = "DefinitionIdentityValue"
+ for tag in ("first", "second"):
+ with self.subTest(tag=tag):
+ runtime = load_type_spec_runtime(
+ compile_type_spec(
+ string_spec(
+ name=type_name,
+ definitions=f'tag(::{type_name}) = "{tag}"',
+ )
+ )
+ )
+ value = runtime.module._sample(runtime.module.Random.default_rng())
+ self.assertEqual(str(runtime.module.tag(value)), tag)
+
+ def test_repeated_runtime_load_does_not_mutate_child_module(self):
+ spec = string_spec(name="RepeatedLoadValue")
+ definition = compile_type_spec_runtime(
+ spec,
+ {1: ["x -> x"]},
+ elementwise_loss="(prediction, target) -> 0.0",
+ loss_function=None,
+ loss_function_expression=None,
+ complexity_mapping=None,
+ early_stop_condition=None,
+ expression_spec="x -> x",
+ expression_spec_function_selector="spec -> spec",
+ )
+ first = load_type_spec_runtime(definition)
+ names_before = tuple(
+ str(name)
+ for name in jl.names(first.configuration_module, all=True, imported=True)
+ )
+
+ second = load_type_spec_runtime(first.definition)
+ names_after = tuple(
+ str(name)
+ for name in jl.names(second.configuration_module, all=True, imported=True)
+ )
+
+ self.assertEqual(names_after, names_before)
+ self.assertTrue(
+ bool(
+ jl.seval("(a, b) -> a === b")(
+ first.expression_spec,
+ second.expression_spec,
+ )
+ )
+ )
+
+ def test_failed_configuration_source_remains_isolated(self):
+ suffix = uuid.uuid4().hex
+ type_name = f"FailedConfigurationValue_{suffix}"
+ operator_name = f"isolated_operator_{suffix}"
+ X, y = string_data()
+ fitted = tiny_model(
+ string_spec(name=type_name),
+ operators={1: [f"{operator_name}(x::{type_name}) = x"]},
+ )
+ fitted.fit(X, y)
+ broken = tiny_model(
+ string_spec(name=type_name),
+ operators={
+ 1: [
+ f"""
+ begin
+ {operator_name}(x::{type_name}) = {type_name}("a")
+ error("isolated source failure")
+ end
+ """
+ ]
+ },
+ )
+
+ with self.assertRaisesRegex(JuliaError, "isolated source failure"):
+ broken.fit(X, y)
+
+ np.testing.assert_array_equal(fitted.predict(X), y)
+ self.assertFalse(bool(jl.isdefined(jl.Main, jl.Symbol(operator_name))))
+
+ def test_failed_runtime_install_can_retry_same_definition(self):
+ suffix = uuid.uuid4().hex
+ flag = f"_pysr_retry_install_{suffix}"
+ type_name = f"RetryInstallValue_{suffix}"
+ operator_name = f"retry_identity_{suffix}"
+ jl.seval(f"global {flag} = Ref(true)")
+ definition = compile_type_spec_runtime(
+ string_spec(name=type_name),
+ {1: [f"{operator_name}(x::{type_name})::{type_name} = x"]},
+ elementwise_loss=f"(x::{type_name}, y::{type_name}) -> 0.0",
+ loss_function=None,
+ loss_function_expression=None,
+ complexity_mapping=None,
+ early_stop_condition=f"""
+ begin
+ Main.{flag}[] && error("retry source failure")
+ (loss, complexity) -> false
+ end
+ """,
+ )
+
+ with self.assertRaisesRegex(JuliaError, "retry source failure"):
+ load_type_spec_runtime(definition)
+
+ jl.seval(f"{flag}[] = false")
+ runtime = load_type_spec_runtime(definition)
+ self.assertEqual(runtime.operator_names, {1: (operator_name,)})
+ self.assertIsNotNone(runtime.early_stop_condition)
+
+ def test_failed_type_install_can_retry_same_definition(self):
+ suffix = uuid.uuid4().hex
+ flag = f"_pysr_retry_type_install_{suffix}"
+ type_name = f"RetryTypeValue_{suffix}"
+ operator_name = f"retry_type_identity_{suffix}"
+ jl.seval(f"global {flag} = Ref(true)")
+ definition = compile_type_spec_runtime(
+ string_spec(
+ name=type_name,
+ preamble=f'Main.{flag}[] && error("retry type failure")',
+ ),
+ {1: [f"{operator_name}(x::{type_name})::{type_name} = x"]},
+ elementwise_loss=f"(x::{type_name}, y::{type_name}) -> 0.0",
+ loss_function=None,
+ loss_function_expression=None,
+ complexity_mapping=None,
+ early_stop_condition=None,
+ )
+
+ with self.assertRaisesRegex(JuliaError, "retry type failure"):
+ load_type_spec_runtime(definition)
+
+ jl.seval(f"{flag}[] = false")
+ runtime = load_type_spec_runtime(definition)
+ self.assertEqual(runtime.operator_names, {1: (operator_name,)})
+ self.assertTrue(
+ bool(jl.isdefined(runtime.module, jl.Symbol(type_name))),
+ )
+
+ def test_equivalent_runtime_definition_is_evaluated_once(self):
+ suffix = uuid.uuid4().hex
+ counter = f"_pysr_definition_count_{suffix}"
+ type_name = f"ReusedValue_{suffix}"
+ operator = f"reused_identity_{suffix}"
+ jl.seval(f"global {counter} = 0")
+ source = f"""
+ begin
+ Main.{counter} += 1
+ {operator}(x::{type_name}) = x
+ end
+ """
+ first = tiny_model(string_spec(name=type_name), operators={1: [source]})
+ second = tiny_model(string_spec(name=type_name), operators={1: [source]})
+ X, y = string_data()
+
+ first.fit(X, y)
+ second.fit(X, y)
+ restored = pickle.loads(pickle.dumps(second))
+ _ = restored.julia_state_
+ _ = restored.julia_options_
+
+ self.assertEqual(int(jl.seval(counter)), 1)
+ self.assertTrue(
+ bool(
+ jl.seval("(a, b) -> a === b")(
+ first._load_type_spec_runtime().configuration_module,
+ second._load_type_spec_runtime().configuration_module,
+ )
+ )
+ )
+
+ def test_restored_model_rebuilds_callable_columns(self):
+ X, y = string_data()
+ model = tiny_model(string_spec())
+ model.fit(X, y)
+
+ restored = pickle.loads(pickle.dumps(model))
+ self.assertNotIn("lambda_format", restored.equations_.columns)
+
+ best = restored.get_best()
+ self.assertIn("lambda_format", restored.equations_.columns)
+ np.testing.assert_array_equal(best["lambda_format"](X), model.predict(X))
+
+ def test_failed_warm_start_restores_fitted_metadata_and_checkpoint(self):
+ X, y = string_data()
+ with tempfile.TemporaryDirectory() as directory:
+ model = tiny_model(
+ string_spec(),
+ temp_equation_file=False,
+ output_directory=directory,
+ run_id="warm-start-rollback",
+ )
+ model.fit(X, y)
+ definition = model._type_spec_runtime_definition_
+ state = model.julia_state_stream_.copy()
+ equations = model.equations_.copy(deep=True)
+ feature_names = model.feature_names_in_.copy()
+ checkpoint = Path(model.get_pkl_filename())
+ checkpoint_contents = checkpoint.read_bytes()
+
+ model.set_params(
+ warm_start=True,
+ complexity_mapping="expression -> 1",
+ )
+ with self.assertRaisesRegex(
+ ValueError, "Cannot warm-start after changing.*configuration"
+ ):
+ model.fit(X, y, variable_names=["candidate"])
+
+ self.assertEqual(model._type_spec_runtime_definition_, definition)
+ np.testing.assert_array_equal(model.julia_state_stream_, state)
+ pd.testing.assert_frame_equal(model.equations_, equations)
+ np.testing.assert_array_equal(model.feature_names_in_, feature_names)
+ self.assertEqual(checkpoint.read_bytes(), checkpoint_contents)
+
+ def fail_after_mutation(*args, **kwargs):
+ model.equations_.iloc[0, 0] += 100
+ model.feature_names_in_[0] = "mutated"
+ raise RuntimeError("search failed")
+
+ model.set_params(complexity_mapping=None)
+ with (
+ patch.object(model, "_run", side_effect=fail_after_mutation),
+ self.assertRaisesRegex(RuntimeError, "search failed"),
+ ):
+ model.fit(X, y, variable_names=["candidate"])
+
+ self.assertEqual(model._type_spec_runtime_definition_, definition)
+ np.testing.assert_array_equal(model.julia_state_stream_, state)
+ pd.testing.assert_frame_equal(model.equations_, equations)
+ np.testing.assert_array_equal(model.feature_names_in_, feature_names)
+ self.assertEqual(checkpoint.read_bytes(), checkpoint_contents)
+
+ def test_warm_start_rejects_loss_type_change(self):
+ X, y = string_data()
+ type_name = "WarmStartLossTypeValue"
+ objective = """
+ function warm_start_objective(
+ tree, dataset::Dataset, options, idx=nothing
+ )::Float64
+ _, complete = eval_tree_array(tree, dataset.X, options)
+ return complete ? 0.0 : Inf
+ end
+ """
+ model = tiny_model(
+ string_spec(name=type_name, loss_type="Float64"),
+ elementwise_loss=None,
+ loss_function=objective,
+ )
+ model.fit(X, y)
+ model.set_params(
+ warm_start=True,
+ type_spec=string_spec(name=type_name, loss_type="Float32"),
+ )
+
+ with self.assertRaisesRegex(
+ ValueError, "Cannot warm-start after changing.*configuration"
+ ):
+ model.fit(X, y)
+
+ def test_pre_search_checkpoint_can_resume_as_a_fresh_fit(self):
+ X, y = string_data()
+ with tempfile.TemporaryDirectory() as directory:
+ model = tiny_model(
+ string_spec(name="PreSearchCheckpointValue"),
+ temp_equation_file=False,
+ output_directory=directory,
+ run_id="pre-search-checkpoint",
+ )
+ with (
+ patch.object(
+ PySRRegressor,
+ "_run",
+ side_effect=RuntimeError("search failed"),
+ ),
+ self.assertRaisesRegex(RuntimeError, "search failed"),
+ ):
+ model.fit(X, y)
+
+ checkpoint = Path(model.get_pkl_filename())
+ with checkpoint.open("rb") as checkpoint_file:
+ restored = pickle.load(checkpoint_file)
+ self.assertIsNone(restored.equations_)
+ self.assertFalse(hasattr(restored, "_type_spec_runtime_definition_"))
+
+ restored.set_params(
+ temp_equation_file=True,
+ output_directory=None,
+ run_id=None,
+ )
+ restored.fit(X, y)
+ np.testing.assert_array_equal(restored.predict(X), y)
+
+ def test_callable_supervised_loss_is_supported(self):
+ type_name = "CallableLossValue"
+ loss_name = "CallableTypeSpecLoss"
+ X, y = string_data()
+ model = tiny_model(
+ string_spec(name=type_name),
+ elementwise_loss=f"""
+ begin
+ struct {loss_name} <: SymbolicRegression.SupervisedLoss end
+ (::{loss_name})(
+ prediction::{type_name}, target::{type_name}
+ )::Float64 = prediction == target ? 0.0 : 1.0
+ {loss_name}()
+ end
+ """,
+ )
+
+ model.fit(X, y)
+ np.testing.assert_array_equal(model.predict(X), y)
+
+ def test_numeric_early_stop_condition_remains_numeric(self):
+ model = tiny_model(
+ string_spec(name="NumericEarlyStopValue"),
+ early_stop_condition=1e-4,
+ )
+ definition = model._compile_type_spec_runtime(model._operators_from_params())
+ runtime = load_type_spec_runtime(definition)
+
+ self.assertEqual(float(runtime.early_stop_condition), 1e-4)
+
+ def test_pickle_restores_runtime_before_deserializing_in_fresh_process(self):
+ X, y = string_data()
+ with tempfile.TemporaryDirectory() as directory:
+ model = tiny_model(
+ string_spec(name="CheckpointStringValue"),
+ complexity_mapping="checkpoint_complexity(expression) = 1",
+ early_stop_condition=(
+ "checkpoint_early_stop(loss, complexity) = false"
+ ),
+ temp_equation_file=False,
+ output_directory=directory,
+ run_id="typespec-checkpoint",
+ delete_tempfiles=False,
+ )
+ model.fit(X, y)
+ type_module_name = (
+ model._type_spec_runtime_definition_.type_spec.module_name
+ )
+ run_directory = Path(directory) / "typespec-checkpoint"
+ module_name = model._type_spec_runtime_definition_.module_name
+ code = f"""
+import json
+import numpy as np
+from pysr import PySRRegressor, jl
+type_module_name = {type_module_name!r}
+module_name = {module_name!r}
+assert not bool(jl.isdefined(jl.Main, jl.Symbol(type_module_name)))
+model = PySRRegressor.from_file(run_directory={str(run_directory)!r})
+X = np.array([["a"], ["b"], ["a"], ["b"]], dtype=object)
+runtime = model._load_type_spec_runtime()
+print(json.dumps({{
+ "prediction": model.predict(X).tolist(),
+ "loaded": bool(jl.isdefined(runtime.module, jl.Symbol(module_name))),
+ "complexity": str(jl.Base.nameof(runtime.complexity_mapping)),
+ "early_stop": str(jl.Base.nameof(runtime.early_stop_condition)),
+}}))
+"""
+ result = subprocess.run(
+ [sys.executable, "-c", code],
+ check=False,
+ capture_output=True,
+ text=True,
+ )
+
+ self.assertEqual(result.returncode, 0, result.stderr)
+ payload = json.loads(result.stdout.strip().splitlines()[-1])
+ self.assertEqual(
+ payload,
+ {
+ "prediction": y.tolist(),
+ "loaded": True,
+ "complexity": "checkpoint_complexity",
+ "early_stop": "checkpoint_early_stop",
+ },
+ )
+
+ def test_unsupported_configurations_fail_before_runtime_loading(self):
+ X, y = string_data()
+ cases = (
+ (tiny_model(string_spec(), turbo=True), ValueError, "turbo"),
+ (tiny_model(string_spec(), bumper=True), ValueError, "bumper"),
+ (
+ tiny_model(string_spec(), autodiff_backend="Zygote"),
+ ValueError,
+ "autodiff_backend",
+ ),
+ (
+ tiny_model(string_spec(), output_jax_format=True),
+ ValueError,
+ "output_jax_format",
+ ),
+ (
+ tiny_model(string_spec(), output_torch_format=True),
+ ValueError,
+ "output_torch_format",
+ ),
+ (
+ tiny_model(string_spec(), extra_sympy_mappings={"f": lambda x: x}),
+ ValueError,
+ "extra_sympy_mappings",
+ ),
+ (
+ tiny_model(
+ string_spec(), extra_jax_mappings={lambda x: x: "jnp.asarray"}
+ ),
+ ValueError,
+ "extra_jax_mappings",
+ ),
+ (
+ tiny_model(
+ string_spec(), extra_torch_mappings={lambda x: x: lambda x: x}
+ ),
+ ValueError,
+ "extra_torch_mappings",
+ ),
+ (
+ tiny_model(string_spec(), binary_operators=["+"]),
+ ValueError,
+ "binary_operators",
+ ),
+ )
+ for model, error, message in cases:
+ with self.subTest(message=message):
+ with self.assertRaisesRegex(error, message):
+ model.fit(X, y)
+ self.assertFalse(hasattr(model, "_type_spec_runtime_definition_"))
+
+ def test_custom_expression_spec_requires_runtime_source(self):
+ class MissingSourceSpec(ExpressionSpec):
+ pass
+
+ X, y = string_data()
+ model = tiny_model(
+ string_spec(name="MissingExpressionSourceValue"),
+ expression_spec=MissingSourceSpec(),
+ )
+
+ with self.assertRaisesRegex(ValueError, "does not provide Julia source"):
+ model.fit(X, y)
+ self.assertFalse(hasattr(model, "_type_spec_runtime_definition_"))
+
+ def test_operator_names_drive_constraint_keys(self):
+ type_name = "NamedOperatorValue"
+ operator_name = "named_identity"
+ X, y = string_data()
+ model = tiny_model(
+ string_spec(name=type_name),
+ operators={1: [f"{operator_name}(x::{type_name})::{type_name} = x"]},
+ constraints={operator_name: 1},
+ )
+
+ model.fit(X, y)
+
+ runtime = model._load_type_spec_runtime()
+ self.assertEqual(runtime.operator_names, {1: (operator_name,)})
+
+ def test_template_source_hook_runs_once_per_fit(self):
+ sources = []
+
+ class SourceExpressionSpec(AbstractExpressionSpec):
+ @property
+ def supports_type_spec(self):
+ return True
+
+ def julia_expression_spec(self):
+ return ExpressionSpec().julia_expression_spec()
+
+ def create_exports(self, *args, **kwargs):
+ return ExpressionSpec().create_exports(*args, **kwargs)
+
+ def _julia_expression_spec_source(self, *, prototype=None):
+ sources.append(prototype)
+ return f"""
+ @template_spec(
+ expressions=(f,), prototype={prototype}
+ ) do x
+ f(x)
+ end
+ """
+
+ def _julia_expression_spec_function_selector(self):
+ return "spec -> spec.structure.combine"
+
+ X, y = string_data()
+ model = tiny_model(
+ string_spec(name="SingleTemplateSourceValue"),
+ expression_spec=SourceExpressionSpec(),
+ )
+ model.fit(X, y)
+
+ self.assertEqual(
+ sources,
+ ["SymbolicRegression.init_value(SingleTemplateSourceValue)"],
+ )
+ np.testing.assert_array_equal(model.predict(X), y)
+
+ def test_template_type_spec_parallelism_matrix(self):
+ type_name = "TemplateParallelismValue"
+ for parallelism in ("serial", "multithreading", "multiprocessing"):
+ X, y = string_data(constant=True)
+ model = tiny_model(
+ string_spec(
+ name=type_name,
+ sample=f'rng -> {type_name}("a")',
+ mutate="(rng, value, temperature) -> value",
+ ),
+ expression_spec=identity_template(),
+ parallelism=parallelism,
+ procs=2 if parallelism == "multiprocessing" else None,
+ niterations=100,
+ early_stop_condition=1e-8,
+ )
+
+ with self.subTest(parallelism=parallelism):
+ model.fit(X, y)
+ np.testing.assert_array_equal(model.predict(X), y)
+
+ def test_same_runtime_warm_start_continues(self):
+ X, y = string_data()
+ model = tiny_model(
+ string_spec(),
+ expression_spec=identity_template(),
+ )
+ model.fit(X, y)
+ fingerprint = model._type_spec_runtime_definition_.fingerprint
+ model.set_params(warm_start=True)
+
+ model.fit(X, y)
+
+ self.assertEqual(
+ model._type_spec_runtime_definition_.fingerprint,
+ fingerprint,
+ )
+ np.testing.assert_array_equal(model.predict(X), y)
+
+ def test_malformed_hooks_fail_with_hook_names(self):
+ cases = (
+ (
+ vector_spec(
+ name="BadScalarConstantsValue",
+ scalar_constants="value -> 1.0",
+ ),
+ "scalar_constants.*AbstractVector",
+ ),
+ (
+ string_spec(
+ name="BadPredicateValue",
+ is_valid="value -> 1",
+ ),
+ "is_valid.*Bool",
+ ),
+ (
+ string_spec(name="BadInitValue", init='() -> ""'),
+ "init.*BadInitValue",
+ ),
+ (
+ vector_spec(
+ name="BadReconstructionValue",
+ with_scalar_constants="(value, constants) -> constants",
+ ),
+ "with_scalar_constants.*BadReconstructionValue",
+ ),
+ (
+ vector_spec(
+ name="BadScalarConstantElementValue",
+ scalar_constants="value -> [1, 2]",
+ with_scalar_constants=(
+ "(value, constants) -> "
+ "BadScalarConstantElementValue(collect(constants))"
+ ),
+ ),
+ "scalar_constants.*concrete `AbstractFloat`",
+ ),
+ (
+ vector_spec(
+ name="NonRoundTripValue",
+ with_scalar_constants=(
+ "(value, constants) -> "
+ "NonRoundTripValue(reverse(collect(constants)))"
+ ),
+ ),
+ "with_scalar_constants.*preserve",
+ ),
+ (
+ vector_spec(name="BadStringValue", string="value -> 1"),
+ "string.*AbstractString",
+ ),
+ (
+ string_spec(
+ name="BadSampleArityValue",
+ sample='() -> BadSampleArityValue("a")',
+ ),
+ "sample",
+ ),
+ (
+ string_spec(
+ name="RejectedSampleValue",
+ is_valid="value -> false",
+ ),
+ "sample.*invalid value",
+ ),
+ )
+ for spec, message in cases:
+ with self.subTest(spec=spec.name):
+ runtime = load_type_spec_runtime(compile_type_spec(spec))
+ with self.assertRaisesRegex(ValueError, message):
+ validate_type_spec_runtime(runtime)
+
+ def test_first_fit_validates_type_spec_hooks(self):
+ X, y = string_data()
+ model = tiny_model(string_spec(name="FitValidationValue", init='() -> ""'))
+
+ with self.assertRaisesRegex(ValueError, "init.*FitValidationValue"):
+ model.fit(X, y)
+ self.assertFalse(hasattr(model, "_type_spec_runtime_definition_"))
+
+ def test_operator_and_elementwise_loss_are_invoked_during_validation(self):
+ type_name = "RuntimeInvocationValue"
+ cases = {
+ "operator": {
+ "operator": (
+ f"function exploding_operator(x::{type_name})::{type_name}\n"
+ ' x.data == "b" && error("operator was invoked")\n'
+ " return x\nend"
+ ),
+ "loss": (
+ f"identity_loss(x::{type_name}, y::{type_name})::Float64 = 0.0"
+ ),
+ "message": "operator was invoked",
+ },
+ "elementwise loss": {
+ "operator": (f"identity_runtime(x::{type_name})::{type_name} = x"),
+ "loss": (
+ f"function exploding_loss(x::{type_name}, y::{type_name})::Float64\n"
+ ' x.data == y.data && error("loss was invoked")\n'
+ " return 0.0\nend"
+ ),
+ "message": "loss was invoked",
+ },
+ }
+ for name, case in cases.items():
+ with self.subTest(case=name):
+ definition = compile_type_spec_runtime(
+ string_spec(name=type_name),
+ {1: [case["operator"]]},
+ elementwise_loss=case["loss"],
+ loss_function=None,
+ loss_function_expression=None,
+ complexity_mapping=None,
+ early_stop_condition=None,
+ )
+ runtime = load_type_spec_runtime(definition)
+ with self.assertRaisesRegex(ValueError, case["message"]):
+ validate_type_spec_runtime(runtime)
+
+ def test_module_loading_has_no_behavioral_hook_side_effects(self):
+ suffix = uuid.uuid4().hex
+ init_count = f"_pysr_init_count_{suffix}"
+ mutate_count = f"_pysr_mutate_count_{suffix}"
+ type_name = f"SideEffectValue_{suffix}"
+ jl.seval(f"global {init_count} = 0; global {mutate_count} = 0")
+ spec = string_spec(
+ name=type_name,
+ init=(f"() -> begin Main.{init_count} += 1; " f'{type_name}("a") end'),
+ mutate=(
+ f"(rng, value, temperature) -> begin Main.{mutate_count} += 1; "
+ "value end"
+ ),
+ )
+
+ load_type_spec_runtime(compile_type_spec(spec))
+ self.assertEqual(int(jl.seval(init_count)), 0)
+ self.assertEqual(int(jl.seval(mutate_count)), 0)
+
+ definition = compile_type_spec_runtime(
+ spec,
+ {1: [f"side_effect_identity(x::{type_name}) = x"]},
+ elementwise_loss=(
+ f"side_effect_loss(x::{type_name}, y::{type_name})::Float64 = 0.0"
+ ),
+ loss_function=None,
+ loss_function_expression=None,
+ complexity_mapping=None,
+ early_stop_condition=None,
+ )
+ runtime = load_type_spec_runtime(definition)
+ validate_type_spec_options(
+ runtime,
+ runtime.operator_functions,
+ runtime.elementwise_loss,
+ )
+ self.assertEqual(int(jl.seval(init_count)), 0)
+ self.assertEqual(int(jl.seval(mutate_count)), 0)
+ validate_type_spec_runtime(runtime)
+ self.assertEqual(int(jl.seval(init_count)), 1)
+ self.assertEqual(int(jl.seval(mutate_count)), 1)
+
+ def test_object_array_shapes_and_logical_conversion(self):
+ one_dimensional = object_array_1d(value for value in ([1, 2], [3, 4]))
+ self.assertEqual(one_dimensional.shape, (2,))
+ self.assertEqual(one_dimensional.tolist(), [[1, 2], [3, 4]])
+ two_dimensional = object_array_2d(
+ row for row in (([1, 2], [3, 4]), ([5, 6], [7, 8]))
+ )
+ self.assertEqual(two_dimensional.shape, (2, 2))
+ with self.assertRaisesRegex(ValueError, "2D array"):
+ object_array_2d(["a", "b"])
+ with self.assertRaisesRegex(ValueError, "same number"):
+ object_array_2d([[1], [2, 3]])
+
+ spec = TypeSpec(
+ "PairValue",
+ fields={"number": "Float64", "label": "String"},
+ sample='rng -> PairValue(0.0, "")',
+ mutate="(rng, value, temperature) -> value",
+ )
+ runtime = load_type_spec_runtime(compile_type_spec(spec))
+ values = object_array_1d([(1.0, "one"), (2.0, "two")])
+ converted = type_spec_to_julia_array(runtime, values)
+ self.assertEqual(
+ type_spec_to_python_array(runtime, converted).tolist(), values.tolist()
+ )
+
+ def test_derived_scalar_constant_packing_honors_offsets(self):
+ runtime = load_type_spec_runtime(
+ compile_type_spec(vector_spec(name="OffsetVectorValue"))
+ )
+ first = runtime.module._convert_value([1.0, 2.0])
+ second = runtime.module._convert_value([3.0, 4.0])
+ packed, next_index, unpacked = jl.seval("""
+ function (first, second)
+ DE = SymbolicRegression.InterfaceDynamicExpressionsModule.DE
+ packed = fill(-99.0, 8)
+ idx = DE.pack_scalar_constants!(packed, 3, first)
+ idx = DE.pack_scalar_constants!(packed, idx, second)
+ idx, rebuilt_first =
+ DE.unpack_scalar_constants(packed, 3, first)
+ idx, rebuilt_second =
+ DE.unpack_scalar_constants(packed, idx, second)
+ return packed, idx, (rebuilt_first, rebuilt_second)
+ end
+ """)(first, second)
+ self.assertEqual(
+ list(packed),
+ [-99.0, -99.0, 1.0, 2.0, 3.0, 4.0, -99.0, -99.0],
+ )
+ self.assertEqual(int(next_index), 7)
+ self.assertEqual(
+ [list(value.data) for value in unpacked],
+ [[1.0, 2.0], [3.0, 4.0]],
+ )
+
+ def test_union_payload_scalar_constant_hooks(self):
+ spec = TypeSpec(
+ "TensorValue",
+ fields={"data": "Union{Float64, Vector{Float64}, Matrix{Float64}}"},
+ sample="rng -> TensorValue(randn(rng, 2, 2))",
+ scalar_constants="""
+ function scalar_constants(value)
+ return value.data isa Float64 ? [value.data] : vec(value.data)
+ end
+ """,
+ with_scalar_constants="""
+ function with_scalar_constants(value, scalar_constants)
+ data = value.data isa Float64 ? scalar_constants[1] :
+ reshape(collect(scalar_constants), size(value.data))
+ return TensorValue(data)
+ end
+ """,
+ )
+ runtime = load_type_spec_runtime(compile_type_spec(spec))
+ value = runtime.module._convert_value(np.eye(2))
+ self.assertEqual(
+ list(runtime.module._scalar_constants(value)),
+ [1.0, 0.0, 0.0, 1.0],
+ )
+
+ def test_immutable_scalar_constant_view_validates(self):
+ spec = vector_spec(
+ name="ImmutableScalarConstantsValue",
+ scalar_constants=(
+ "value -> range(value.data[1], value.data[end]; "
+ "length=length(value.data))"
+ ),
+ with_scalar_constants=(
+ "(value, scalar_constants) -> "
+ "ImmutableScalarConstantsValue(collect(scalar_constants))"
+ ),
+ )
+ model = tiny_model(spec)
+ runtime = model._load_type_spec_runtime(
+ for_fit=True,
+ operators=model._operators_from_params(),
+ )
+ self.assertEqual(runtime.spec.name, "ImmutableScalarConstantsValue")
+
+ def test_variable_length_vector_constants(self):
+ spec = TypeSpec(
+ "VariableVector",
+ fields={"data": "Vector{Float64}"},
+ sample="rng -> VariableVector(randn(rng, rand(rng, 1:5)))",
+ mutate="""
+ function mutate_vector(rng, value, temperature)
+ if rand(rng) < 0.2
+ return VariableVector(randn(rng, rand(rng, 1:5)))
+ end
+ data = copy(value.data)
+ data[rand(rng, eachindex(data))] += temperature * randn(rng)
+ return VariableVector(data)
+ end
+ """,
+ scalar_constants="value -> value.data",
+ with_scalar_constants=(
+ "(value, scalar_constants) -> "
+ "VariableVector(collect(scalar_constants))"
+ ),
+ )
+ rng = np.random.default_rng(0)
+ values = [rng.normal(size=2) for _ in range(64)]
+ X = pd.DataFrame({"x": values})
+ prefix = np.array([1.5])
+ suffix = np.array([-0.2, -3.0, 0.1])
+ y = np.empty(len(values), dtype=object)
+ y[:] = [np.concatenate((prefix, value, suffix)) for value in values]
+ model = tiny_model(
+ spec,
+ expression_spec=identity_template(),
+ operators={
+ 2: ["concat_vectors(a, b) = " "VariableVector(vcat(a.data, b.data))"]
+ },
+ elementwise_loss="""
+ function vector_loss(a, b)::Float64
+ return length(a.data) == length(b.data) ?
+ sum(abs2, a.data - b.data) : 1.0e6
+ end
+ """,
+ niterations=60,
+ ncycles_per_iteration=100,
+ populations=4,
+ population_size=50,
+ tournament_selection_n=10,
+ maxsize=7,
+ early_stop_condition=(
+ "(loss, complexity) -> loss < 1.0e-8 && complexity == 5"
+ ),
+ should_optimize_constants=True,
+ )
+
+ model.fit(X, y)
+
+ exact = model.equations_.query("complexity == 5").sort_values("loss").iloc[0]
+ self.assertLess(exact.loss, 1.0e-8)
+ constants = jl.seval("""
+ function (expression)
+ DE = SymbolicRegression.InterfaceDynamicExpressionsModule.DE
+ tree = DE.get_tree(expression)
+ nodes = DE.filter_map(
+ node -> node.degree == 0 && node.constant,
+ identity,
+ tree,
+ typeof(tree),
+ )
+ return map(node -> node.val.data, nodes)
+ end
+ """)(exact.julia_expression)
+ self.assertEqual([len(value) for value in constants], [1, 3])
+ np.testing.assert_allclose(constants[0], prefix, atol=1.0e-6)
+ np.testing.assert_allclose(constants[1], suffix, atol=1.0e-6)
+
+ def test_multi_field_fit_predicts_tuples(self):
+ spec = TypeSpec(
+ "FitPairValue",
+ fields={"number": "Float64", "label": "String"},
+ sample='rng -> FitPairValue(0.0, "")',
+ mutate="(rng, value, temperature) -> value",
+ )
+ pairs = [(1.0, "one"), (2.0, "two"), (3.0, "three"), (4.0, "four")]
+ X = pd.DataFrame({"x": pairs})
+ y = pd.Series(pairs, dtype=object)
+ model = tiny_model(spec)
+
+ model.fit(X, y)
+
+ self.assertEqual(model.predict(X).tolist(), pairs)
+
+ def test_invalid_expression_evaluation_reports_the_rejected_value(self):
+ type_name = "InvalidEvaluationValue"
+ spec = TypeSpec(
+ type_name,
+ fields={"data": "Float64"},
+ sample=f"rng -> {type_name}(1.0)",
+ scalar_constants="value -> (value.data,)",
+ with_scalar_constants=f"(value, constants) -> {type_name}(constants[1])",
+ )
+ definition = compile_type_spec_runtime(
+ spec,
+ {1: [f"invalid_operator(x::{type_name})::{type_name} = {type_name}(Inf)"]},
+ elementwise_loss="(prediction, target) -> 0.0",
+ loss_function=None,
+ loss_function_expression=None,
+ complexity_mapping=None,
+ early_stop_condition=None,
+ )
+ runtime = load_type_spec_runtime(definition)
+ expression = jl.seval("""
+ (value_type, operator) -> begin
+ operators = OperatorEnum(1 => (operator,))
+ node = Node{value_type}(op=1, l=Node{value_type}(feature=1))
+ Expression(node; operators, variable_names=["x1"])
+ end
+ """)(runtime.value_type, runtime.operator_functions[1][0])
+ X = object_array_2d([[1.0], [2.0]])
+
+ with self.assertRaisesRegex(ValueError, "`is_valid` rejected"):
+ CallableJuliaExpression(expression, runtime)(X)
+
+ def test_type_spec_variable_name_count_is_validated(self):
+ X, y = string_data()
+ model = tiny_model(string_spec(name="WrongVariableCountValue"))
+
+ with self.assertRaisesRegex(
+ ValueError, "variable_names.*one name per TypeSpec feature"
+ ):
+ model.fit(X, y, variable_names=["first", "extra"])
+
+ def test_rejects_empty_feature_axis(self):
+ model = tiny_model(string_spec())
+
+ with self.assertRaisesRegex(ValueError, "at least one feature"):
+ model.fit(np.empty((2, 0), dtype=object), np.array(["a", "b"]))
+
+ def test_csv_only_loading_is_rejected(self):
+ with tempfile.TemporaryDirectory() as directory:
+ Path(directory, "hall_of_fame.csv").write_text(
+ "Complexity,Loss,Equation\n1,0.0,x0\n",
+ encoding="utf-8",
+ )
+ with self.assertRaisesRegex(ValueError, "checkpoint.pkl"):
+ PySRRegressor.from_file(
+ run_directory=directory,
+ type_spec=string_spec(),
+ operators={1: ["identity_value(x::StringValue) = x"]},
+ n_features_in=1,
+ )
+
+ def test_template_custom_combiner_infers_num_features(self):
+ type_name = "TemplateVectorValue"
+ model = tiny_model(
+ vector_spec(name=type_name),
+ expression_spec=TemplateExpressionSpec(
+ combine="add_vectors(f(x1), g(x2))",
+ expressions=["f", "g"],
+ variable_names=["x1", "x2"],
+ ),
+ operators={
+ 2: [
+ f"""
+ add_vectors(a::{type_name}, b::{type_name}) =
+ {type_name}(a.data + b.data)
+ add_vectors(a::ValidVector, b::ValidVector) =
+ ValidVector(map(add_vectors, a.x, b.x), a.valid && b.valid)
+ """
+ ]
+ },
+ elementwise_loss=(
+ f"vector_loss(a::{type_name}, b::{type_name})::Float64 = "
+ "sum(abs2, a.data - b.data)"
+ ),
+ )
+ X = np.empty((4, 2), dtype=object)
+ X[:, 0] = [np.array([1.0, 2.0])] * 4
+ X[:, 1] = [np.array([3.0, 4.0])] * 4
+ y = np.empty(4, dtype=object)
+ y[:] = [np.array([4.0, 6.0])] * 4
+
+ model.fit(X, y, variable_names=["x1", "x2"])
+
+ structure = model.julia_options_.expression_options.structure
+ self.assertEqual(int(structure.num_features.f), 1)
+ self.assertEqual(int(structure.num_features.g), 1)
+
+ def test_template_type_spec_parameters(self):
+ type_name = "TemplateParameterValue"
+ X, y = string_data(constant=True)
+ expression_spec = TemplateExpressionSpec(
+ combine="choose_parameter(p[1], f(x))",
+ expressions=["f"],
+ variable_names=["x"],
+ parameters={"p": 1},
+ )
+ model = tiny_model(
+ string_spec(
+ name=type_name,
+ sample=f'rng -> {type_name}("a")',
+ mutate="(rng, value, temperature) -> value",
+ ),
+ expression_spec=expression_spec,
+ operators={
+ 1: [f"identity_{type_name}(x::{type_name}) = x"],
+ 2: [
+ f"""
+ choose_parameter(a::{type_name}, b::{type_name}) = a
+ choose_parameter(a::{type_name}, b::ValidVector) =
+ ValidVector(map(_ -> a, b.x), b.valid)
+ """
+ ],
+ },
+ parallelism="multiprocessing",
+ procs=2,
+ )
+
+ model.fit(X, y)
+ fingerprint = model._type_spec_runtime_definition_.fingerprint
+ model.set_params(warm_start=True)
+
+ model.fit(X, y)
+
+ self.assertEqual(model._type_spec_runtime_definition_.fingerprint, fingerprint)
+ np.testing.assert_array_equal(model.predict(X), y)
+
+ def test_template_type_spec_parameters_are_mutated(self):
+ type_name = "MutatedParameterValue"
+ X = np.array([["a"], ["b"], ["a"], ["b"]], dtype=object)
+ y = np.full(len(X), "b", dtype=object)
+ model = tiny_model(
+ string_spec(
+ name=type_name,
+ sample=f'rng -> {type_name}("a")',
+ mutate=f'(rng, value, temperature) -> {type_name}("b")',
+ ),
+ expression_spec=TemplateExpressionSpec(
+ combine="choose_parameter(p[1], f(x))",
+ expressions=["f"],
+ variable_names=["x"],
+ parameters={"p": 1},
+ ),
+ operators={
+ 1: [f"identity_{type_name}(x::{type_name}) = x"],
+ 2: [
+ f"""
+ choose_parameter(a::{type_name}, b::{type_name}) = a
+ choose_parameter(a::{type_name}, b::ValidVector) =
+ ValidVector(map(_ -> a, b.x), b.valid)
+ """
+ ],
+ },
+ niterations=10,
+ ncycles_per_iteration=10,
+ populations=2,
+ population_size=20,
+ )
+
+ model.fit(X, y)
+
+ np.testing.assert_array_equal(model.predict(X), y)
+ self.assertEqual(model.get_best()["loss"], 0.0)
+
+ def test_numeric_path_remains_independent(self):
+ X = np.linspace(-1.0, 1.0, 20).reshape(-1, 1)
+ y = X[:, 0]
+ model = PySRRegressor(
+ unary_operators=[],
+ binary_operators=["+"],
+ niterations=1,
+ ncycles_per_iteration=2,
+ populations=1,
+ population_size=8,
+ tournament_selection_n=3,
+ maxsize=7,
+ parallelism="serial",
+ deterministic=True,
+ random_state=0,
+ progress=False,
+ verbosity=0,
+ temp_equation_file=True,
+ should_optimize_constants=False,
+ )
+ model.fit(X, y)
+ prediction = model.predict(X)
+ self.assertEqual(prediction.shape, y.shape)
+ self.assertTrue(np.isfinite(prediction).all())
+
+ model.set_params(
+ type_spec=string_spec(),
+ operators={1: ["identity_value(x::StringValue) = x"]},
+ )
+
+ np.testing.assert_array_equal(model.predict(X), prediction)
+ self.assertTrue(model._supports_export("sympy"))
+
+ def test_configured_full_objective_and_module_local_helpers(self):
+ X, y = string_data()
+ type_name = "ConfiguredObjectiveValue"
+ called = "CONFIGURED_OBJECTIVE_SAW_BATCH"
+ model = tiny_model(
+ string_spec(
+ name=type_name,
+ preamble=(
+ f"const {called} = Ref(false)\n" "const PrivateLoss = Float64"
+ ),
+ loss_type="PrivateLoss",
+ ),
+ operators={1: [f"configured_identity(x::{type_name}) = x"]},
+ elementwise_loss=None,
+ loss_function="""
+ function configured_objective(
+ tree, dataset::Dataset, options, idx=nothing
+ )::PrivateLoss
+ CONFIGURED_OBJECTIVE_SAW_BATCH[] |= idx !== nothing
+ _, complete = eval_tree_array(tree, dataset.X, options)
+ return complete ? 0.0 : Inf
+ end
+ """,
+ batching=True,
+ batch_size=2,
+ complexity_mapping=f"expression -> {type_name} <: Any ? 1 : 1",
+ early_stop_condition="(loss, complexity) -> false",
+ )
+ model.fit(X, y)
+ runtime = model._load_type_spec_runtime()
+
+ self.assertTrue(
+ bool(
+ jl.getindex(
+ jl.getproperty(runtime.configuration_module, jl.Symbol(called))
+ )
+ )
+ )
+ self.assertEqual(model.predict(X).shape, y.shape)
+
+ def test_custom_expression_specs_are_deferred(self):
+ class CustomExpressionSpec(AbstractExpressionSpec):
+ def julia_expression_spec(self):
+ return ExpressionSpec().julia_expression_spec()
+
+ def create_exports(self, *args, **kwargs):
+ return ExpressionSpec().create_exports(*args, **kwargs)
+
+ X, y = string_data()
+ model = tiny_model(
+ string_spec(name="DeferredExpressionValue"),
+ expression_spec=CustomExpressionSpec(),
+ )
+ with self.assertRaisesRegex(
+ ValueError, "CustomExpressionSpec does not support TypeSpec"
+ ):
+ model.fit(X, y)
+
+ def test_rejects_invalid_specs(self):
+ cases = [
+ (dict(name="not an id"), "not an identifier"),
+ (dict(fields={}), "non-empty ordered mapping"),
+ (dict(fields={"not an id": "String"}), "is not an identifier"),
+ (dict(fields={"data": " "}), "requires a Julia type"),
+ (dict(sample=" "), "must contain Julia source"),
+ (dict(scalar_constants="value -> [1.0]"), "must be provided together"),
+ (dict(mutate=None), "requires an explicit `mutate`"),
+ (dict(string=" "), "cannot be empty"),
+ ]
+ for overrides, message in cases:
+ with self.subTest(**overrides):
+ with self.assertRaisesRegex(ValueError, message):
+ string_spec(**overrides)
+
+ def test_rejects_invalid_configurations(self):
+ no_loss = dict(
+ elementwise_loss=None, loss_function=None, loss_function_expression=None
+ )
+ operators = {1: ["x -> x"]}
+ cases = [
+ (string_spec(), operators, no_loss, "exactly one of"),
+ (
+ string_spec(loss_type="Float64"),
+ operators,
+ {**no_loss, "elementwise_loss": "(p, t) -> 0.0"},
+ "return type is inferred",
+ ),
+ (
+ string_spec(),
+ operators,
+ {**no_loss, "loss_function": "f(tree, dataset, options) = 0.0"},
+ "explicit `loss_type`",
+ ),
+ (string_spec(), None, no_loss, "requires explicit"),
+ (string_spec(), {0: ["x -> x"]}, no_loss, "positive integers"),
+ (string_spec(), {1: []}, no_loss, "cannot be empty"),
+ (string_spec(), {1: [" "]}, no_loss, "must contain Julia source"),
+ ]
+ for spec, operator_table, losses, message in cases:
+ with self.subTest(message=message):
+ with self.assertRaisesRegex(ValueError, message):
+ validate_type_spec_configuration(spec, operator_table, **losses)
+
+ def test_fit_data_validation(self):
+ X = np.array([["a"], ["b"]], dtype=object)
+ y = X[:, 0].copy()
+
+ def prepare(
+ model=None,
+ X=X,
+ y=y,
+ Xresampled=None,
+ weights=None,
+ variable_names=None,
+ X_units=None,
+ y_units=None,
+ ):
+ return prepare_type_spec_fit_data(
+ model or PySRRegressor(),
+ X,
+ y,
+ Xresampled,
+ weights,
+ variable_names,
+ None,
+ X_units,
+ y_units,
+ )
+
+ with self.assertRaisesRegex(NotImplementedError, "denoising"):
+ prepare(model=PySRRegressor(denoise=True))
+ with self.assertRaisesRegex(ValueError, "2D array"):
+ prepare(X=np.empty((1, 1, 1), dtype=object))
+ with self.assertRaisesRegex(NotImplementedError, "one output"):
+ prepare(y=np.empty((2, 2), dtype=object))
+ with self.assertRaisesRegex(ValueError, "inconsistent numbers of samples"):
+ prepare(y=y[:1])
+ with self.assertRaisesRegex(ValueError, "at least one sample"):
+ prepare(X=np.empty((0, 1), dtype=object), y=np.empty(0, dtype=object))
+ with self.assertRaisesRegex(NotImplementedError, "weights"):
+ prepare(weights=np.ones(2))
+ with self.assertRaisesRegex(NotImplementedError, "units"):
+ prepare(X_units=["m"])
+ with self.assertWarnsRegex(UserWarning, "reset to `None`"):
+ prepare(X=pd.DataFrame({"a": ["a", "b"]}), variable_names=["a"])
+ model = PySRRegressor()
+ with self.assertWarnsRegex(UserWarning, "Spaces"):
+ prepare(model=model, y=y.reshape(-1, 1), variable_names=["a b"])
+ self.assertEqual(list(model.feature_names_in_), ["a_b"])
+
+ def test_prediction_data_validation(self):
+ model = PySRRegressor()
+ model.n_features_in_ = 1
+ model.selection_mask_ = None
+ model.feature_names_in_ = np.array(["a_b"])
+ with self.assertRaisesRegex(ValueError, "2D array"):
+ prepare_type_spec_prediction_data(model, np.empty((1, 1, 1), dtype=object))
+ with self.assertRaisesRegex(ValueError, "different number of features"):
+ prepare_type_spec_prediction_data(
+ model, np.array([["a", "b"]], dtype=object)
+ )
+ with self.assertRaisesRegex(ValueError, "missing features"):
+ prepare_type_spec_prediction_data(model, pd.DataFrame({"z": ["a"]}))
+ with self.assertWarnsRegex(UserWarning, "Spaces"):
+ out = prepare_type_spec_prediction_data(model, pd.DataFrame({"a b": ["a"]}))
+ self.assertEqual(out.shape, (1, 1))
+
+ model.n_features_in_ = 2
+ model.selection_mask_ = np.array([True, False])
+ out = prepare_type_spec_prediction_data(
+ model, np.array([["a", "b"]], dtype=object)
+ )
+ np.testing.assert_array_equal(out, np.array([["a"]], dtype=object))
+
+ def test_julia_array_conversion_errors(self):
+ runtime = load_type_spec_runtime(
+ compile_type_spec(string_spec(name="ConversionErrorValue"))
+ )
+ with self.assertRaisesRegex(ValueError, "1D or 2D"):
+ type_spec_to_julia_array(runtime, np.empty((1, 1, 1), dtype=object))
+ with self.assertRaises(ValueError):
+ type_spec_to_julia_array(runtime, np.array([1.0, 2.0]))
+
+ def test_runtime_load_rejects_inconsistent_definition_values(self):
+ definition = compile_type_spec_runtime(
+ string_spec(name="InconsistentLoadValue"),
+ {1: ["x -> x"]},
+ elementwise_loss="(prediction, target) -> 0.0",
+ loss_function=None,
+ loss_function_expression=None,
+ complexity_mapping=None,
+ early_stop_condition=None,
+ )
+ load_type_spec_runtime(definition)
+ with patch("pysr.type_specs._runtime_sources", return_value=["x -> x"] * 99):
+ with self.assertRaisesRegex(RuntimeError, "inconsistent"):
+ load_type_spec_runtime(definition)
+
+ def test_fitted_model_guards(self):
+ X, y = string_data()
+ spec = string_spec()
+ model = tiny_model(spec)
+ model.fit(X, y)
+
+ with self.assertRaises(NotImplementedError):
+ model.score(X, y)
+ with self.assertRaises(ValueError):
+ model.predict(np.array([[1.0]]))
+
+ model.set_params(warm_start=True, type_spec=None)
+ with self.assertRaisesRegex(ValueError, "Cannot warm-start"):
+ model.fit(X, y)
+
+ model.set_params(warm_start=False, type_spec=spec)
+ with warnings.catch_warnings():
+ warnings.simplefilter("ignore")
+ with self.assertRaisesRegex(NotImplementedError, "weights"):
+ model.fit(X, y, weights=np.ones(len(y)))
+ self.assertFalse(hasattr(model, "_type_spec_runtime_definition_"))
+
+ def test_guess_seeds_the_search_with_a_custom_constant(self):
+ type_name = "GuessVectorValue"
+ model = tiny_model(
+ vector_spec(name=type_name),
+ operators={
+ 2: [
+ f"add_vectors(a::{type_name}, b::{type_name}) = "
+ f"{type_name}(a.data + b.data)"
+ ]
+ },
+ elementwise_loss=(
+ f"vector_loss(a::{type_name}, b::{type_name})::Float64 = "
+ "sum(abs2, a.data - b.data)"
+ ),
+ guesses=[f"add_vectors(x0, {type_name}([1.5, -2.5]))"],
+ )
+ X = np.empty((4, 1), dtype=object)
+ X[:, 0] = [np.array([1.0, 2.0])] * 4
+ y = np.empty(4, dtype=object)
+ y[:] = [np.array([2.5, -0.5])] * 4
+
+ model.fit(X, y)
+
+ self.assertEqual(model.equations_.loss.min(), 0.0)
+ self.assertIn("[1.5, -2.5]", " ".join(model.equations_.equation))
+
+ def test_template_guess_seeds_each_expression(self):
+ type_name = "GuessTemplateValue"
+ model = tiny_model(
+ string_spec(name=type_name),
+ expression_spec=TemplateExpressionSpec(
+ combine="f(x)",
+ expressions=["f"],
+ variable_names=["x"],
+ ),
+ operators={
+ 2: [
+ f"concat(a::{type_name}, b::{type_name}) = "
+ f"{type_name}(a.data * b.data)"
+ ]
+ },
+ elementwise_loss=(
+ f"string_loss(a::{type_name}, b::{type_name})::Float64 = "
+ "a.data == b.data ? 0.0 : 1.0"
+ ),
+ guesses=[{"f": f'concat(#1, {type_name}("!"))'}],
+ )
+ X = np.array([["a"], ["b"], ["a"], ["b"]], dtype=object)
+ y = np.array(["a!", "b!", "a!", "b!"], dtype=object)
+
+ model.fit(X, y, variable_names=["x"])
+
+ self.assertEqual(model.equations_.loss.min(), 0.0)
+ self.assertIn('concat(#1, "!")', " ".join(model.equations_.equation))
+
+ def test_template_parameter_guess_with_string_values(self):
+ type_name = "GuessTemplateParameterString"
+ model = tiny_model(
+ string_spec(
+ name=type_name,
+ sample=f'rng -> {type_name}("a")',
+ mutate="(rng, value, temperature) -> value",
+ ),
+ expression_spec=TemplateExpressionSpec(
+ combine="choose_parameter(p[1], f(x))",
+ expressions=["f"],
+ variable_names=["x"],
+ parameters={"p": 1},
+ ),
+ operators={
+ 1: [f"identity_{type_name}(x::{type_name}) = x"],
+ 2: [
+ f"""
+ choose_parameter(a::{type_name}, b::{type_name}) = a
+ choose_parameter(a::{type_name}, b::ValidVector) =
+ ValidVector(map(_ -> a, b.x), b.valid)
+ """
+ ],
+ },
+ guesses=[{"f": "#1", "p": ["!"]}],
+ )
+ X = np.array([["a"], ["b"], ["a"], ["b"]], dtype=object)
+ y = np.full(len(X), "!", dtype=object)
+
+ model.fit(X, y, variable_names=["x"])
+
+ np.testing.assert_array_equal(model.predict(X), y)
+
+ def test_template_parameter_guess_with_vector_values(self):
+ type_name = "GuessTemplateParameterVector"
+ parameter_values = object_array_1d([np.array([1.5, -2.5])])
+ original_value = parameter_values[0].copy()
+ model = tiny_model(
+ vector_spec(name=type_name),
+ expression_spec=TemplateExpressionSpec(
+ combine="choose_parameter(p[1], f(x))",
+ expressions=["f"],
+ variable_names=["x"],
+ parameters={"p": 1},
+ ),
+ operators={
+ 1: [f"identity_{type_name}(x::{type_name}) = x"],
+ 2: [
+ f"""
+ choose_parameter(a::{type_name}, b::{type_name}) = a
+ choose_parameter(a::{type_name}, b::ValidVector) =
+ ValidVector(map(_ -> a, b.x), b.valid)
+ """
+ ],
+ },
+ guesses=[{"f": "#1", "p": parameter_values}],
+ )
+ X = np.empty((4, 1), dtype=object)
+ X[:, 0] = [np.array([1.0, 2.0])] * 4
+ y = np.empty(4, dtype=object)
+ y[:] = [np.array([1.5, -2.5])] * 4
+
+ model.fit(X, y, variable_names=["x"])
+
+ self.assertEqual(
+ [list(value) for value in model.predict(X)],
+ [list(value) for value in y],
+ )
+ np.testing.assert_array_equal(parameter_values[0], original_value)
+
+ def test_template_parameter_guess_with_packet_values(self):
+ type_name = "GuessTemplateParameterPacket"
+ model = tiny_model(
+ TypeSpec(
+ type_name,
+ fields={"number": "Float64", "label": "String"},
+ sample=f'rng -> {type_name}(0.0, "")',
+ mutate="(rng, value, temperature) -> value",
+ ),
+ expression_spec=TemplateExpressionSpec(
+ combine="choose_parameter(p[1], f(x))",
+ expressions=["f"],
+ variable_names=["x"],
+ parameters={"p": 1},
+ ),
+ operators={
+ 1: [f"identity_{type_name}(x::{type_name}) = x"],
+ 2: [
+ f"""
+ choose_parameter(a::{type_name}, b::{type_name}) = a
+ choose_parameter(a::{type_name}, b::ValidVector) =
+ ValidVector(map(_ -> a, b.x), b.valid)
+ """
+ ],
+ },
+ guesses=[{"f": "#1", "p": [(2.0, "two")]}],
+ )
+ X = np.empty((4, 1), dtype=object)
+ X[:, 0] = object_array_1d([(1.0, "one")] * 4)
+ y = object_array_1d([(2.0, "two")] * 4)
+
+ model.fit(X, y, variable_names=["x"])
+
+ self.assertEqual(model.predict(X).tolist(), y.tolist())
+
+ def test_guess_rejects_a_constant_that_is_not_the_custom_type(self):
+ type_name = "GuessConversionValue"
+ model = tiny_model(
+ vector_spec(name=type_name),
+ guesses=[f"identity_{type_name}(2.0)"],
+ )
+ X = np.empty((4, 1), dtype=object)
+ X[:, 0] = [np.array([1.0, 2.0])] * 4
+ y = np.empty(4, dtype=object)
+ y[:] = [np.array([1.0, 2.0])] * 4
+
+ with self.assertRaisesRegex(
+ JuliaError, f"Cannot `convert` an object of type Float64 to .*{type_name}"
+ ):
+ model.fit(X, y)
+
+ def test_guess_rejects_an_undefined_constant_constructor(self):
+ model = tiny_model(
+ string_spec(),
+ guesses=['MissingGuessValue("!")'],
+ )
+ X, y = string_data()
+
+ with self.assertRaisesRegex(
+ JuliaError, "Tried to interpolate function `MissingGuessValue`"
+ ):
+ model.fit(X, y)
+
+
+if __name__ == "__main__":
+ unittest.main()
diff --git a/pysr/type_specs.py b/pysr/type_specs.py
new file mode 100644
index 000000000..351b16229
--- /dev/null
+++ b/pysr/type_specs.py
@@ -0,0 +1,1368 @@
+from __future__ import annotations
+
+import copy
+import hashlib
+import json
+import warnings
+from dataclasses import dataclass
+from functools import cache
+from textwrap import dedent
+from typing import Any
+
+import numpy as np
+import pandas as pd
+from juliacall import JuliaError # type: ignore
+
+from .julia_helpers import jl_array
+from .julia_import import AnyValue, SymbolicRegression, jl
+
+_TYPE_MODULE_INSTALLED = "_PYSR_TYPE_SPEC_INSTALLED"
+
+
+def object_array_1d(values: Any) -> np.ndarray:
+ """Build a 1D object array whose cells are the logical values themselves."""
+ if isinstance(values, np.ndarray):
+ return values if values.dtype == object else values.astype(object)
+ values = list(values)
+ array = np.empty(len(values), dtype=object)
+ for i, value in enumerate(values):
+ array[i] = value
+ return array
+
+
+def object_array_2d(values: Any) -> np.ndarray:
+ """Build a 2D object array whose cells are the logical values themselves."""
+ if isinstance(values, np.ndarray):
+ return values if values.dtype == object else values.astype(object)
+ try:
+ values = list(values)
+ if any(isinstance(row, (str, bytes)) for row in values):
+ raise TypeError
+ rows = [list(row) for row in values]
+ except TypeError:
+ raise ValueError("TypeSpec X must be a 2D array of logical values.")
+ n_columns = len(rows[0]) if rows else 0
+ array = np.empty((len(rows), n_columns), dtype=object)
+ for i, row in enumerate(rows):
+ if len(row) != n_columns:
+ raise ValueError("All rows of X must have the same number of features.")
+ for j, value in enumerate(row):
+ array[i, j] = value
+ return array
+
+
+@dataclass
+class TypeSpec:
+ """Definition of a custom value type for symbolic regression.
+
+ Parameters
+ ----------
+ name : str
+ Name of the generated Julia type. Use this name in operator and hook
+ definitions.
+ fields : dict[str, str]
+ Ordered mapping from field names to Julia field types.
+ sample : str
+ Julia callable with signature ``rng -> (...)::{name}``.
+ scalar_constants : str, optional
+ Julia callable with signature ``value::{name} -> (...)::Vector{Float64}``. Provide
+ together with ``with_scalar_constants`` to enable continuous constant
+ optimization.
+ with_scalar_constants : str, optional
+ Julia callable with signature ``(value::{name}, scalar_constants::Vector{Float64}) -> (...)::{name}``.
+ init : str, optional
+ Julia callable with signature ``() -> value::{name}``. The default repurposes `sample`
+ using a deterministic local random-number generator.
+ mutate : str, optional
+ Julia callable with signature ``(rng, value::{name}, temperature::Float64) -> value::{name}``.
+ Use this to mutate a value of the type.
+ The `temperature` (from simulated annealing) is between 0 and 1. High temperatures
+ generally should be used for more aggressive mutations, and low temperatures for
+ more conservative mutations.
+ The default mutation repurposes `scalar_constants` and `with_scalar_constants`
+ to mutate the set with SymbolicRegression.jl's default constant mutation,
+ and put it back into a new value.
+ is_valid : str, optional
+ Julia callable with signature ``value::{name} -> (...)::Bool``.
+ Use this to check if your type is in a valid state. For example, if your type
+ has been processed in a way that invalidates it (like a NaN, but specific to your type),
+ you may use this function to check for that. This will be used to quit evaluation early.
+ The default checks that every scalar constant is finite, or accepts every value for a
+ non-optimizable type.
+ init_invalid : str, optional
+ Julia callable with signature ``() -> value::{name}`` that constructs an
+ invalid value. The result must have this type and fail ``is_valid``.
+ Used by DynamicExpressions to fill failed convenience evaluations.
+ Omit this hook for types with no invalid state. It receives no prototype
+ or shape information.
+ string : str, optional
+ Julia callable with signature ``value::{name} -> (...)::AbstractString`` used to print
+ values.
+ preamble : str, optional
+ Julia source evaluated once before the generated type definition. Types
+ and functions it defines are visible to the hooks and to operator and
+ objective sources; imports are not, so any source needing a package
+ beyond SymbolicRegression must import it itself.
+ definitions : str, optional
+ Julia source evaluated once directly after the generated type, for
+ constructors and methods that mention it.
+ loss_type : str, optional
+ Concrete Julia ``AbstractFloat`` type returned by a custom full
+ objective. Elementwise loss return types are inferred.
+ """
+
+ name: str
+ fields: dict[str, str]
+ sample: str
+ scalar_constants: str | None = None
+ with_scalar_constants: str | None = None
+ init: str | None = None
+ mutate: str | None = None
+ is_valid: str | None = None
+ string: str | None = None
+ preamble: str | None = None
+ definitions: str | None = None
+ loss_type: str | None = None
+ init_invalid: str | None = None
+
+ def __post_init__(self) -> None:
+ if not isinstance(self.name, str) or not self.name.isidentifier():
+ raise ValueError(f"TypeSpec name {self.name!r} is not an identifier.")
+ if not isinstance(self.fields, dict) or not self.fields:
+ raise ValueError("`fields` must be a non-empty ordered mapping.")
+ for name, field_type in self.fields.items():
+ if not isinstance(name, str) or not name.isidentifier():
+ raise ValueError(f"TypeSpec field name {name!r} is not an identifier.")
+ if not isinstance(field_type, str) or not field_type.strip():
+ raise ValueError(f"TypeSpec field `{name}` requires a Julia type.")
+ if not isinstance(self.sample, str) or not self.sample.strip():
+ raise ValueError("`sample` must contain Julia source.")
+
+ if (self.scalar_constants is None) != (self.with_scalar_constants is None):
+ raise ValueError(
+ "`scalar_constants` and `with_scalar_constants` must be provided "
+ "together."
+ )
+ if not self.can_optimize and self.mutate is None:
+ raise ValueError(
+ "A non-optimizable TypeSpec requires an explicit `mutate` callable."
+ )
+
+ for name in (
+ "scalar_constants",
+ "with_scalar_constants",
+ "init",
+ "mutate",
+ "is_valid",
+ "init_invalid",
+ "string",
+ "preamble",
+ "definitions",
+ "loss_type",
+ ):
+ value = getattr(self, name)
+ if value is not None and (not isinstance(value, str) or not value.strip()):
+ raise ValueError(f"`{name}` cannot be empty.")
+
+ @property
+ def can_optimize(self) -> bool:
+ return self.scalar_constants is not None
+
+
+@dataclass(frozen=True)
+class _TypeSpecDefinition:
+ spec: TypeSpec
+ fingerprint: str
+ source: str
+
+ @property
+ def module_name(self) -> str:
+ return f"_PySRTypeSpec_{self.fingerprint[:20]}"
+
+
+@dataclass(frozen=True)
+class _TypeSpecRuntimeDefinition:
+ type_spec: _TypeSpecDefinition
+ operators: tuple[tuple[int, tuple[str, ...]], ...]
+ elementwise_loss: str | None
+ loss_function: str | None
+ loss_function_expression: str | None
+ complexity_mapping: str | None
+ early_stop_condition: str | None
+ expression_spec: str
+ expression_spec_function_selector: str | None
+ fingerprint: str
+ source: str
+
+ @property
+ def module_name(self) -> str:
+ return f"_PySRConfig_{self.fingerprint[:20]}"
+
+
+@dataclass(frozen=True)
+class _TypeSpecRuntime:
+ definition: _TypeSpecDefinition | _TypeSpecRuntimeDefinition
+ module: AnyValue
+ value_type: AnyValue
+ configuration_module: AnyValue
+ operator_functions: dict[int, tuple[AnyValue, ...]]
+ operator_names: dict[int, tuple[str, ...]]
+ elementwise_loss: AnyValue | None = None
+ loss_function: AnyValue | None = None
+ loss_function_expression: AnyValue | None = None
+ complexity_mapping: AnyValue | None = None
+ early_stop_condition: float | AnyValue | None = None
+ expression_spec: AnyValue | None = None
+
+ @property
+ def type_definition(self) -> _TypeSpecDefinition:
+ if isinstance(self.definition, _TypeSpecRuntimeDefinition):
+ return self.definition.type_spec
+ return self.definition
+
+ @property
+ def spec(self) -> TypeSpec:
+ return self.type_definition.spec
+
+
+def _quoted(source: str) -> str:
+ return json.dumps(source, ensure_ascii=False).replace("$", r"\$")
+
+
+def _block(source: str) -> str:
+ return dedent(source).strip()
+
+
+def _runtime_sources(
+ operators: tuple[tuple[int, tuple[str, ...]], ...],
+ elementwise_loss: str | None,
+ loss_function: str | None,
+ loss_function_expression: str | None,
+ complexity_mapping: str | None,
+ early_stop_condition: str | None,
+ expression_spec: str,
+) -> tuple[str, ...]:
+ sources = tuple(
+ source for _, operator_sources in operators for source in operator_sources
+ )
+ optional_sources = (
+ elementwise_loss,
+ loss_function,
+ loss_function_expression,
+ complexity_mapping,
+ early_stop_condition,
+ )
+ return (
+ *sources,
+ *(source for source in optional_sources if source is not None),
+ expression_spec,
+ )
+
+
+def _runtime_install_source(
+ type_definition: _TypeSpecDefinition,
+ runtime_module_name: str,
+ sources: tuple[str, ...],
+ expression_spec_function_selector: str | None,
+) -> str:
+ source_literals = ",\n".join(f" {_quoted(source)}" for source in sources)
+ selector = _optional_source(expression_spec_function_selector)
+ return _block(f"""
+ let type_module_name = Symbol({_quoted(type_definition.module_name)})
+ installed =
+ isdefined(Main, type_module_name) &&
+ isdefined(
+ Base.invokelatest(getproperty, Main, type_module_name),
+ Symbol({_quoted(_TYPE_MODULE_INSTALLED)}),
+ )
+ if !installed
+ Base.include_string(
+ Main,
+ {_quoted(type_definition.source)},
+ {_quoted("PySR." + type_definition.module_name)},
+ )
+ end
+ end
+ let
+ parent = getproperty(
+ Main, Symbol({_quoted(type_definition.module_name)})
+ )
+ module_name = Symbol({_quoted(runtime_module_name)})
+ needs_install = !isdefined(parent, module_name)
+ if !needs_install
+ existing = Base.invokelatest(getproperty, parent, module_name)
+ needs_install = !isdefined(existing, :_definition_values)
+ end
+ if needs_install
+ Core.eval(parent, Expr(:module, true, module_name, Expr(:block)))
+ module_ = Base.invokelatest(getproperty, parent, module_name)
+ Core.eval(module_, :(using SymbolicRegression))
+ for name in getproperty(parent, :_PYSR_PARENT_BINDING_NAMES)
+ isdefined(module_, name) && continue
+ Core.eval(
+ module_,
+ Expr(
+ :const,
+ Expr(
+ :(=),
+ name,
+ QuoteNode(getproperty(parent, name)),
+ ),
+ ),
+ )
+ end
+ sources = String[
+ {source_literals}
+ ]
+ values = Any[]
+ for (index, source) in enumerate(sources)
+ push!(
+ values,
+ Base.include_string(
+ module_,
+ source,
+ "PySR TypeSpec configuration $index",
+ ),
+ )
+ end
+ selector_source = {selector}
+ if selector_source !== nothing
+ selector_value = Base.include_string(
+ module_,
+ selector_source,
+ "PySR TypeSpec expression function selector",
+ )
+ push!(values, Base.invokelatest(selector_value, values[end]))
+ end
+ Core.eval(
+ module_,
+ Expr(
+ :const,
+ Expr(:(=), :_definition_values, QuoteNode(Tuple(values))),
+ ),
+ )
+ end
+ end
+ """) + "\n"
+
+
+def compile_type_spec_runtime(
+ spec: TypeSpec,
+ operators: dict[int, list[str]],
+ *,
+ elementwise_loss: str | None,
+ loss_function: str | None,
+ loss_function_expression: str | None,
+ complexity_mapping: str | None,
+ early_stop_condition: str | None,
+ expression_spec: str | None = None,
+ expression_spec_function_selector: str | None = None,
+) -> _TypeSpecRuntimeDefinition:
+ validate_type_spec_configuration(
+ spec,
+ operators,
+ elementwise_loss=elementwise_loss,
+ loss_function=loss_function,
+ loss_function_expression=loss_function_expression,
+ )
+ type_definition = compile_type_spec(spec)
+ normalized_operators = tuple(
+ (arity, tuple(sources)) for arity, sources in _normalize_operators(operators)
+ )
+ if expression_spec is None:
+ expression_spec = _block(f"""
+ SymbolicRegression.ExpressionSpec(
+ node_type=SymbolicRegression.InterfaceDynamicExpressionsModule.DE.Node{{
+ {spec.name}
+ }},
+ )
+ """)
+ payload = {
+ "type_spec": type_definition.fingerprint,
+ "operators": normalized_operators,
+ "elementwise_loss": elementwise_loss,
+ "loss_function": loss_function,
+ "loss_function_expression": loss_function_expression,
+ "complexity_mapping": complexity_mapping,
+ "early_stop_condition": early_stop_condition,
+ "expression_spec": expression_spec,
+ "expression_spec_function_selector": expression_spec_function_selector,
+ "loss_type": spec.loss_type,
+ }
+ fingerprint = hashlib.sha256(
+ json.dumps(payload, sort_keys=True).encode()
+ ).hexdigest()
+ runtime_module_name = f"_PySRConfig_{fingerprint[:20]}"
+ sources = _runtime_sources(
+ normalized_operators,
+ elementwise_loss,
+ loss_function,
+ loss_function_expression,
+ complexity_mapping,
+ early_stop_condition,
+ expression_spec,
+ )
+ source = _runtime_install_source(
+ type_definition,
+ runtime_module_name,
+ sources,
+ expression_spec_function_selector,
+ )
+ return _TypeSpecRuntimeDefinition(
+ type_spec=type_definition,
+ operators=normalized_operators,
+ elementwise_loss=elementwise_loss,
+ loss_function=loss_function,
+ loss_function_expression=loss_function_expression,
+ complexity_mapping=complexity_mapping,
+ early_stop_condition=early_stop_condition,
+ expression_spec=expression_spec,
+ expression_spec_function_selector=expression_spec_function_selector,
+ fingerprint=fingerprint,
+ source=source,
+ )
+
+
+def compile_type_spec_runtime_for_model(
+ model: Any,
+ operators: dict[int, list[str]],
+) -> _TypeSpecRuntimeDefinition:
+ """Compile one model configuration without evaluating its Julia sources."""
+ from .expression_specs import ExpressionSpec
+
+ spec = model.type_spec
+ assert spec is not None
+ expression_spec = model.expression_spec_
+ expression_spec_source = expression_spec._julia_expression_spec_source(
+ prototype=f"SymbolicRegression.init_value({spec.name})"
+ )
+ if expression_spec_source is None and type(expression_spec) is not ExpressionSpec:
+ raise ValueError(
+ f"{type(expression_spec).__name__} declares TypeSpec support "
+ "but does not provide Julia source."
+ )
+ return compile_type_spec_runtime(
+ spec,
+ operators,
+ elementwise_loss=model.elementwise_loss,
+ loss_function=model.loss_function,
+ loss_function_expression=model.loss_function_expression,
+ complexity_mapping=model.complexity_mapping,
+ early_stop_condition=(
+ None
+ if model.early_stop_condition is None
+ else str(model.early_stop_condition)
+ ),
+ expression_spec=expression_spec_source,
+ expression_spec_function_selector=(
+ expression_spec._julia_expression_spec_function_selector()
+ ),
+ )
+
+
+_TYPE_SPEC_MODULE = _block(r"""
+ import SymbolicRegression: init_value, mutate_value, sample_value
+ import SymbolicRegression.ConstantOptimizationModule: can_optimize
+ import SymbolicRegression.InterfaceDynamicExpressionsModule: string_constant
+ import SymbolicRegression.InterfaceDynamicExpressionsModule.DE:
+ count_scalar_constants, get_number_type, is_valid, invalid_value,
+ pack_scalar_constants!, unpack_scalar_constants
+ import SymbolicRegression.InterfaceDynamicExpressionsModule.DE.StringsModule:
+ needs_brackets
+
+ const _config = __TYPE_SPEC_CONFIG__
+ abstract type _TypeSpecValue end
+
+ macro _define_type_spec(config_expression)
+ config = Core.eval(__module__, config_expression)
+ fields = map(config.fields) do (name, type)
+ :($(Symbol(name))::$(Meta.parse(type)))
+ end
+ type_definition = Expr(
+ :struct,
+ false,
+ Expr(:<:, Symbol(config.name), :_TypeSpecValue),
+ Expr(:block, fields...),
+ )
+ return esc(type_definition)
+ end
+
+ _include(source, label) = Base.include_string(@__MODULE__, source, label)
+ _config.preamble === nothing ||
+ _include(_config.preamble, "TypeSpec.preamble")
+ @_define_type_spec _config
+ const _value_type = getfield(@__MODULE__, Symbol(_config.name))
+ _config.definitions === nothing ||
+ _include(_config.definitions, "TypeSpec.definitions")
+ SymbolicRegression.parse_scope(::Type{<:_TypeSpecValue}) = @__MODULE__
+
+ _fields(value) = ntuple(i -> getfield(value, i), fieldcount(_value_type))
+ Base.:(==)(a::_TypeSpecValue, b::_TypeSpecValue) = _fields(a) == _fields(b)
+ Base.isequal(a::_TypeSpecValue, b::_TypeSpecValue) =
+ isequal(_fields(a), _fields(b))
+ Base.hash(value::_TypeSpecValue, h::UInt) = hash(_fields(value), h)
+
+ const _sample = _include(_config.sample, "TypeSpec.sample")
+ const _init = _config.init === nothing ?
+ () -> _sample(Random.Xoshiro(0)) :
+ _include(_config.init, "TypeSpec.init")
+ init_value(::Type{<:_TypeSpecValue}) = _init()
+ sample_value(rng::AbstractRNG, ::Type{<:_TypeSpecValue}, options) = _sample(rng)
+ can_optimize(::Type{<:_TypeSpecValue}, _) = _config.optimizable
+
+ if _config.optimizable
+ const _scalar_constants =
+ _include(_config.scalar_constants, "TypeSpec.scalar_constants")
+ const _with_scalar_constants = _include(
+ _config.with_scalar_constants,
+ "TypeSpec.with_scalar_constants",
+ )
+ get_number_type(::Type{<:_TypeSpecValue}) =
+ eltype(_scalar_constants(init_value(_value_type)))
+ end
+
+ count_scalar_constants(value::_TypeSpecValue) =
+ _config.optimizable ? length(_scalar_constants(value)) : 0
+
+ if _config.optimizable
+ function pack_scalar_constants!(
+ buffer::AbstractVector{<:Number}, idx::Int, value::_TypeSpecValue
+ )
+ scalar_constants = _scalar_constants(value)
+ copyto!(
+ buffer,
+ idx,
+ scalar_constants,
+ firstindex(scalar_constants),
+ length(scalar_constants),
+ )
+ return idx + length(scalar_constants)
+ end
+ function unpack_scalar_constants(
+ buffer::AbstractVector{<:Number}, idx::Int, value::_TypeSpecValue
+ )
+ count = length(_scalar_constants(value))
+ scalar_constants = @view buffer[idx:(idx + count - 1)]
+ return idx + count, _with_scalar_constants(value, scalar_constants)
+ end
+ end
+
+ const _mutate = if _config.mutate === nothing
+ function (rng, value, temperature, mutation)
+ scalar_constants = collect(_scalar_constants(value))
+ isempty(scalar_constants) && return _sample(rng)
+ i = rand(rng, eachindex(scalar_constants))
+ scalar_constants[i] =
+ SymbolicRegression.MutationFunctionsModule.mutate_value(
+ rng, scalar_constants[i], temperature, mutation
+ )
+ return _with_scalar_constants(value, scalar_constants)
+ end
+ else
+ mutate = _include(_config.mutate, "TypeSpec.mutate")
+ (rng, value, temperature, _) -> mutate(rng, value, temperature)
+ end
+ mutate_value(
+ rng::AbstractRNG,
+ value::_TypeSpecValue,
+ temperature,
+ mutation::SymbolicRegression.ConstantMutation,
+ ) = _mutate(rng, value, temperature, mutation)
+
+ const _is_valid = if _config.is_valid !== nothing
+ _include(_config.is_valid, "TypeSpec.is_valid")
+ elseif _config.optimizable
+ value -> all(isfinite, _scalar_constants(value))
+ else
+ _ -> true
+ end
+ is_valid(value::_TypeSpecValue) = _is_valid(value)
+
+ if _config.init_invalid !== nothing
+ const _init_invalid = _include(_config.init_invalid, "TypeSpec.init_invalid")
+ invalid_value(::Type{<:_TypeSpecValue}) = _init_invalid()
+ end
+
+ const _string = if _config.string === nothing
+ function (value)
+ fields = map(1:fieldcount(_value_type)) do i
+ sprint(show, getfield(value, i); context=:compact => true)
+ end
+ return fieldcount(_value_type) == 1 ? only(fields) :
+ string(_config.name, "(", join(fields, ", "), ")")
+ end
+ else
+ _include(_config.string, "TypeSpec.string")
+ end
+ Base.show(io::IO, value::_TypeSpecValue) = print(io, _string(value))
+ needs_brackets(::_TypeSpecValue) = false
+ string_constant(
+ value::_TypeSpecValue, ::Val{precision}, unit
+ ) where {precision} = _string(value) * unit
+
+ function _convert_value(x)
+ x isa _value_type && return x
+ x = PythonCall.Py(x)
+ fieldcount(_value_type) == 1 ||
+ PythonCall.pyhasattr(x, "__len__") &&
+ PythonCall.pylen(x) == fieldcount(_value_type) ||
+ throw(ArgumentError(
+ "TypeSpec values require exactly $(fieldcount(_value_type)) fields."
+ ))
+ values = ntuple(fieldcount(_value_type)) do i
+ source = fieldcount(_value_type) == 1 ? x : x[i - 1]
+ PythonCall.pyconvert(fieldtype(_value_type, i), source)
+ end
+ return _value_type(values...)
+ end
+ _convert_array(values, dims) = reshape(
+ _value_type[_convert_value(value) for value in values], Tuple(dims)
+ )
+ """)
+
+
+def _optional_source(source: str | None) -> str:
+ return "nothing" if source is None else _quoted(source)
+
+
+def validate_type_spec_loss_configuration(
+ spec: TypeSpec,
+ elementwise_loss: str | None,
+ loss_function: str | None,
+ loss_function_expression: str | None,
+) -> str:
+ configured = [
+ ("elementwise_loss", elementwise_loss),
+ ("loss_function", loss_function),
+ ("loss_function_expression", loss_function_expression),
+ ]
+ selected = [mode for mode, source in configured if source is not None]
+ if len(selected) != 1:
+ raise ValueError(
+ "TypeSpec requires exactly one of `elementwise_loss`, `loss_function`, "
+ "and `loss_function_expression`."
+ )
+ mode = selected[0]
+ if mode == "elementwise_loss" and spec.loss_type is not None:
+ raise ValueError(
+ "Do not set `loss_type` with `elementwise_loss`; its return type is inferred."
+ )
+ if mode != "elementwise_loss" and spec.loss_type is None:
+ raise ValueError("TypeSpec full objectives require an explicit `loss_type`.")
+ return mode
+
+
+def _normalize_operators(
+ operators: dict[int, list[str]] | None,
+) -> list[tuple[int, list[str]]]:
+ if not operators:
+ raise ValueError("TypeSpec requires explicit `operators={...}`.")
+ normalized = []
+ for arity, sources in operators.items():
+ if type(arity) is not int or arity < 1:
+ raise ValueError("TypeSpec operator arities must be positive integers.")
+ if not sources:
+ raise ValueError(f"TypeSpec operator arity {arity} cannot be empty.")
+ if any(not isinstance(source, str) or not source.strip() for source in sources):
+ raise ValueError("Every TypeSpec operator must contain Julia source.")
+ normalized.append((arity, list(sources)))
+ return sorted(normalized)
+
+
+def validate_type_spec_configuration(
+ spec: TypeSpec,
+ operators: dict[int, list[str]] | None,
+ *,
+ elementwise_loss: str | None,
+ loss_function: str | None,
+ loss_function_expression: str | None,
+) -> None:
+ _normalize_operators(operators)
+ validate_type_spec_loss_configuration(
+ spec, elementwise_loss, loss_function, loss_function_expression
+ )
+
+
+def compile_type_spec(spec: TypeSpec) -> _TypeSpecDefinition:
+ """Create deterministic Julia source without evaluating user code."""
+ field_sources = ",\n".join(
+ f" {_quoted(name)} => {_quoted(field_type)}"
+ for name, field_type in spec.fields.items()
+ )
+ config = _block(f"""
+ (
+ name = {_quoted(spec.name)},
+ fields = (
+ {field_sources},
+ ),
+ sample = {_quoted(spec.sample)},
+ scalar_constants = {_optional_source(spec.scalar_constants)},
+ with_scalar_constants = {_optional_source(spec.with_scalar_constants)},
+ init = {_optional_source(spec.init)},
+ mutate = {_optional_source(spec.mutate)},
+ is_valid = {_optional_source(spec.is_valid)},
+ init_invalid = {_optional_source(spec.init_invalid)},
+ string = {_optional_source(spec.string)},
+ preamble = {_optional_source(spec.preamble)},
+ definitions = {_optional_source(spec.definitions)},
+ optimizable = {str(spec.can_optimize).lower()},
+ )
+ """)
+ body = _TYPE_SPEC_MODULE.replace("__TYPE_SPEC_CONFIG__", config, 1)
+ fingerprint = hashlib.sha256(body.encode()).hexdigest()
+ module_name = f"_PySRTypeSpec_{fingerprint[:20]}"
+ source = _block(f"""
+ module {module_name}
+ using Random
+ using SymbolicRegression
+ using PythonCall
+
+ {body}
+
+ const _PYSR_PARENT_BINDING_NAMES = Tuple(sort!(
+ filter(
+ name -> isdefined(@__MODULE__, name) &&
+ !startswith(String(name), '#'),
+ collect(names(@__MODULE__; all=true, imported=true)),
+ );
+ by=String,
+ ))
+ const {_TYPE_MODULE_INSTALLED} = true
+ end
+ import .{module_name}: {spec.name}
+ """) + "\n"
+ return _TypeSpecDefinition(
+ spec=copy.deepcopy(spec),
+ fingerprint=fingerprint,
+ source=source,
+ )
+
+
+def load_type_spec_runtime(
+ definition: _TypeSpecDefinition | _TypeSpecRuntimeDefinition,
+) -> _TypeSpecRuntime:
+ """Load deterministic TypeSpec modules and return their Julia objects."""
+ type_definition = (
+ definition.type_spec
+ if isinstance(definition, _TypeSpecRuntimeDefinition)
+ else definition
+ )
+ if isinstance(definition, _TypeSpecRuntimeDefinition):
+ install_source = definition.source
+ filename = "PySR." + definition.module_name
+ else:
+ install_source = type_definition.source
+ filename = "PySR." + type_definition.module_name
+ module_symbol = jl.Symbol(type_definition.module_name)
+ runtime_module_symbol = (
+ jl.Symbol(definition.module_name)
+ if isinstance(definition, _TypeSpecRuntimeDefinition)
+ else None
+ )
+ needs_install = not bool(jl.isdefined(jl.Main, module_symbol))
+ if not needs_install:
+ module = jl.getproperty(jl.Main, module_symbol)
+ needs_install = not bool(
+ jl.isdefined(module, jl.Symbol(_TYPE_MODULE_INSTALLED))
+ )
+ if not needs_install and runtime_module_symbol is not None:
+ needs_install = not bool(jl.isdefined(module, runtime_module_symbol))
+ if not needs_install:
+ configuration_module = jl.getproperty(module, runtime_module_symbol)
+ needs_install = not bool(
+ jl.isdefined(configuration_module, jl.Symbol("_definition_values"))
+ )
+ if needs_install:
+ jl.Base.include_string(jl.Main, install_source, filename)
+ module = jl.getproperty(jl.Main, module_symbol)
+ value_type = jl.getproperty(module, jl.Symbol(type_definition.spec.name))
+
+ if isinstance(definition, _TypeSpecDefinition):
+ runtime = _TypeSpecRuntime(
+ definition=definition,
+ module=module,
+ value_type=value_type,
+ configuration_module=module,
+ operator_names={},
+ operator_functions={},
+ )
+ else:
+ assert runtime_module_symbol is not None
+ configuration_module = jl.getproperty(module, runtime_module_symbol)
+ raw_values = list(
+ jl.getproperty(configuration_module, jl.Symbol("_definition_values"))
+ )
+ source_count = len(
+ _runtime_sources(
+ definition.operators,
+ definition.elementwise_loss,
+ definition.loss_function,
+ definition.loss_function_expression,
+ definition.complexity_mapping,
+ definition.early_stop_condition,
+ definition.expression_spec,
+ )
+ )
+ expected_count = source_count + (
+ definition.expression_spec_function_selector is not None
+ )
+ if len(raw_values) != expected_count:
+ raise RuntimeError(
+ "TypeSpec runtime source/value ordering is inconsistent."
+ )
+ values = iter(raw_values[:source_count])
+ operator_functions = {
+ arity: tuple(next(values) for _ in operator_sources)
+ for arity, operator_sources in definition.operators
+ }
+ optional_sources = (
+ definition.elementwise_loss,
+ definition.loss_function,
+ definition.loss_function_expression,
+ definition.complexity_mapping,
+ definition.early_stop_condition,
+ )
+ optional_values = [
+ next(values) if source is not None else None for source in optional_sources
+ ]
+ runtime = _TypeSpecRuntime(
+ definition=definition,
+ module=module,
+ value_type=value_type,
+ configuration_module=configuration_module,
+ operator_functions=operator_functions,
+ operator_names={
+ arity: tuple(str(jl.Base.nameof(function)) for function in functions)
+ for arity, functions in operator_functions.items()
+ },
+ elementwise_loss=optional_values[0],
+ loss_function=optional_values[1],
+ loss_function_expression=optional_values[2],
+ complexity_mapping=optional_values[3],
+ early_stop_condition=optional_values[4],
+ expression_spec=next(values),
+ )
+ return runtime
+
+
+@cache
+def _type_spec_worker_addprocs_factory() -> AnyValue:
+ return jl.seval(r"""
+ begin
+ import Distributed
+ function (
+ addprocs_function,
+ source,
+ filename,
+ worker_imports,
+ type_module_name,
+ runtime_module_name,
+ )
+ addprocs = something(addprocs_function, Distributed.addprocs)
+ expression = Meta.parseall(source)
+ type_module_symbol = Symbol(type_module_name)
+ runtime_module_symbol = Symbol(runtime_module_name)
+ return function (numprocs; kwargs...)
+ procs = addprocs(numprocs; kwargs...)
+ try
+ SymbolicRegression.import_module_on_workers(
+ procs,
+ pathof(SymbolicRegression),
+ worker_imports,
+ 0,
+ )
+ Distributed.remotecall_eval(Main, procs, expression)
+ head_module = getproperty(
+ getproperty(Main, type_module_symbol),
+ runtime_module_symbol,
+ )
+ module_expression = :(
+ getproperty(
+ getproperty(
+ Main,
+ $(QuoteNode(type_module_symbol)),
+ ),
+ $(QuoteNode(runtime_module_symbol)),
+ )
+ )
+ for proc in procs
+ worker_module = Distributed.remotecall_fetch(
+ Core.eval,
+ proc,
+ Main,
+ module_expression,
+ )
+ worker_module === head_module || error(
+ "TypeSpec runtime identity differs on worker $proc."
+ )
+ end
+ catch
+ Distributed.rmprocs(procs)
+ rethrow()
+ end
+ return procs
+ end
+ end
+ end
+ """)
+
+
+def create_type_spec_addprocs_function(
+ definition: _TypeSpecRuntimeDefinition,
+ addprocs_function: AnyValue | None,
+ worker_imports: AnyValue | None,
+) -> AnyValue:
+ """Wrap worker creation so every new process installs the exact runtime source."""
+ return _type_spec_worker_addprocs_factory()(
+ addprocs_function,
+ definition.source,
+ "PySR." + definition.module_name,
+ worker_imports,
+ definition.type_spec.module_name,
+ definition.module_name,
+ )
+
+
+@cache
+def _type_spec_validator() -> AnyValue:
+ return jl.seval(r"""
+ function (module_, T, type_name, optimizable)
+ DE = SymbolicRegression.InterfaceDynamicExpressionsModule.DE
+ fail(hook, message) = throw(ArgumentError("TypeSpec `$hook` $message"))
+ function call(hook, f, args...)
+ try
+ return f(args...)
+ catch error
+ fail(hook, "failed: $(sprint(showerror, error))")
+ end
+ end
+ function check_value(hook, value)
+ value isa T || fail(hook, "must return `$type_name`.")
+ valid = call("is_valid", DE.is_valid, value)
+ valid isa Bool || fail("is_valid", "must return `Bool`.")
+ valid || fail(hook, "returned an invalid value.")
+ return value
+ end
+ function scalar_constant_count(value)
+ count = call(
+ "count_scalar_constants", DE.count_scalar_constants, value
+ )
+ count isa Int && count >= 0 ||
+ fail("count_scalar_constants", "must return a nonnegative `Int`.")
+ return count
+ end
+ function check_optimization(value, count)
+ scalar_constants =
+ call("scalar_constants", module_._scalar_constants, value)
+ scalar_constants isa AbstractVector ||
+ fail("scalar_constants", "must return an `AbstractVector`.")
+ length(scalar_constants) == count ||
+ fail("count_scalar_constants", "disagrees with `scalar_constants`.")
+ number_type = DE.get_number_type(T)
+ isconcretetype(number_type) && number_type <: AbstractFloat ||
+ fail("scalar_constants", "must return a vector with a concrete `AbstractFloat` element type.")
+
+ offset = 3
+ packed = fill(number_type(NaN), count + 4)
+ next_idx = call(
+ "pack_scalar_constants!",
+ DE.pack_scalar_constants!,
+ packed,
+ offset,
+ value,
+ )
+ next_idx == offset + count ||
+ fail("pack_scalar_constants!", "returned the wrong next index.")
+ all(isnan, packed[1:(offset - 1)]) &&
+ all(isnan, packed[(offset + count):end]) ||
+ fail("pack_scalar_constants!", "wrote outside its scalar-constant range.")
+ isequal(packed[offset:(offset + count - 1)], scalar_constants) ||
+ fail("pack_scalar_constants!", "disagrees with `scalar_constants`.")
+
+ rebuilt = call(
+ "with_scalar_constants",
+ module_._with_scalar_constants,
+ value,
+ scalar_constants,
+ )
+ rebuilt isa T ||
+ fail("with_scalar_constants", "must return `$type_name`.")
+ isequal(rebuilt, value) ||
+ fail("with_scalar_constants(value, scalar_constants(value))", "must preserve the value.")
+
+ result = call(
+ "unpack_scalar_constants",
+ DE.unpack_scalar_constants,
+ packed,
+ offset,
+ value,
+ )
+ result isa Tuple && length(result) == 2 ||
+ fail("unpack_scalar_constants", "must return `(next_idx, $type_name)`.")
+ unpacked_idx, unpacked = result
+ unpacked_idx == offset + count ||
+ fail("unpack_scalar_constants", "returned the wrong next index.")
+ unpacked isa T ||
+ fail("unpack_scalar_constants", "must return `$type_name`.")
+
+ repacked = fill(number_type(NaN), count + 4)
+ repacked_idx = call(
+ "pack_scalar_constants!",
+ DE.pack_scalar_constants!,
+ repacked,
+ offset,
+ unpacked,
+ )
+ repacked_idx == offset + count && isequal(packed, repacked) ||
+ fail("scalar-constant optimization hooks", "must preserve the packed scalar representation.")
+ end
+
+ rng = module_.Random.Xoshiro(0)
+ sampled = check_value("sample", call("sample", SymbolicRegression.sample_value, rng, T, nothing))
+ initial = check_value("init", call("init", SymbolicRegression.init_value, T))
+ for value in (initial, sampled)
+ count = scalar_constant_count(value)
+ optimizable && check_optimization(value, count)
+ end
+
+ if module_._config.init_invalid !== nothing
+ invalid = call("init_invalid", DE.invalid_value, T)
+ invalid isa T || fail("init_invalid", "must return `$type_name`.")
+ valid = call("is_valid", DE.is_valid, invalid)
+ valid isa Bool || fail("is_valid", "must return `Bool`.")
+ !valid || fail("init_invalid", "must return an invalid value.")
+ end
+
+ mutated = check_value(
+ "mutate",
+ call("mutate", SymbolicRegression.mutate_value, rng, sampled, 1.0, SymbolicRegression.ConstantMutation()),
+ )
+ count = scalar_constant_count(mutated)
+ optimizable && check_optimization(mutated, count)
+
+ call("string", module_._string, sampled) isa AbstractString ||
+ fail("string", "must return an `AbstractString`.")
+ return sampled
+ end
+ """)
+
+
+def validate_type_spec_runtime(runtime: _TypeSpecRuntime) -> None:
+ """Invoke TypeSpec hooks once before searching a new configuration."""
+ try:
+ sampled = _type_spec_validator()(
+ runtime.module,
+ runtime.value_type,
+ runtime.spec.name,
+ runtime.spec.can_optimize,
+ )
+ except JuliaError as error:
+ raise ValueError(str(jl.sprint(jl.showerror, error.args[0]))) from error
+ validate_type_spec_options(
+ runtime,
+ runtime.operator_functions,
+ runtime.elementwise_loss,
+ probe_value=sampled,
+ )
+
+
+@cache
+def _type_spec_operator_validator() -> AnyValue:
+ return jl.seval(r"""
+ function (T, type_name, operator, arity, probe_value)
+ inferred = Base.promote_op(operator, ntuple(_ -> T, arity)...)
+ inferred === T || throw(ArgumentError(
+ "TypeSpec operator `$(nameof(operator))` must be type-stable and " *
+ "infer `$type_name` as its return type; inferred `$inferred`. " *
+ "Add an explicit `::$type_name` return annotation if needed."
+ ))
+ if probe_value !== nothing
+ args = ntuple(_ -> probe_value, arity)
+ result = Base.invokelatest(operator, args...)
+ result isa T || throw(ArgumentError(
+ "TypeSpec operator `$(nameof(operator))` must return `$type_name`."
+ ))
+ end
+ return nothing
+ end
+ """)
+
+
+@cache
+def _type_spec_loss_validator() -> AnyValue:
+ return jl.seval(r"""
+ function (T, elementwise_loss, configured_loss_type, probe_value)
+ loss_type = if elementwise_loss === nothing
+ configured_loss_type
+ else
+ Base.promote_op(elementwise_loss, T, T)
+ end
+ isconcretetype(loss_type) && loss_type <: AbstractFloat ||
+ throw(ArgumentError(
+ "The TypeSpec loss must return a concrete subtype of " *
+ "`AbstractFloat`; got `$loss_type`. Add a concrete Julia " *
+ "return type annotation."
+ ))
+ if probe_value !== nothing && elementwise_loss !== nothing
+ loss = Base.invokelatest(
+ elementwise_loss, probe_value, probe_value
+ )
+ loss isa loss_type || throw(ArgumentError(
+ "The TypeSpec elementwise loss must return `$loss_type`."
+ ))
+ end
+ return loss_type
+ end
+ """)
+
+
+def validate_type_spec_options(
+ runtime: _TypeSpecRuntime,
+ operators: dict[int, tuple[AnyValue, ...]],
+ elementwise_loss: AnyValue | None,
+ *,
+ probe_value: AnyValue | None = None,
+) -> AnyValue:
+ """Validate ordinary Julia options against a loaded TypeSpec."""
+ try:
+ for arity, functions in operators.items():
+ for function in functions:
+ _type_spec_operator_validator()(
+ runtime.value_type,
+ runtime.spec.name,
+ function,
+ arity,
+ probe_value,
+ )
+ configured_loss_type = (
+ None
+ if runtime.spec.loss_type is None
+ else jl.Base.include_string(
+ runtime.module,
+ runtime.spec.loss_type,
+ "PySR TypeSpec loss_type",
+ )
+ )
+ return _type_spec_loss_validator()(
+ runtime.value_type,
+ elementwise_loss,
+ configured_loss_type,
+ probe_value,
+ )
+ except JuliaError as error:
+ raise ValueError(str(jl.sprint(jl.showerror, error.args[0]))) from error
+
+
+def type_spec_to_julia_array(
+ runtime: _TypeSpecRuntime, values: Any, *, transpose: bool = False
+) -> AnyValue:
+ """Convert logical Python values to an array of the generated Julia type."""
+ array = (
+ values if isinstance(values, np.ndarray) else np.asarray(values, dtype=object)
+ )
+ if array.dtype != object:
+ array = array.astype(object)
+ if transpose:
+ array = array.T
+ if array.ndim not in (1, 2):
+ raise ValueError("TypeSpec data must be a 1D or 2D array.")
+
+ try:
+ return runtime.module._convert_array(array.ravel(order="F"), array.shape)
+ except JuliaError as error:
+ raise ValueError(str(jl.sprint(jl.showerror, error.args[0]))) from error
+
+
+def type_spec_to_python_array(runtime: _TypeSpecRuntime, values: Any) -> np.ndarray:
+ """Unwrap generated values into their logical one-field or tuple payloads."""
+ field_names = tuple(runtime.spec.fields)
+ output = np.empty(len(values), dtype=object)
+ for i, value in enumerate(values):
+ if len(field_names) == 1:
+ output[i] = getattr(value, field_names[0])
+ else:
+ output[i] = tuple(getattr(value, name) for name in field_names)
+ return output
+
+
+class CallableJuliaExpression:
+ def __init__(self, expression: AnyValue, runtime: _TypeSpecRuntime):
+ self.expression = expression
+ self.runtime = runtime
+
+ def __call__(self, X: np.ndarray, *args):
+ jl_X = type_spec_to_julia_array(
+ self.runtime,
+ object_array_2d(X),
+ transpose=True,
+ )
+ raw_output, completed = SymbolicRegression.eval_tree_array(
+ self.expression, jl_X, *args
+ )
+ if not bool(completed):
+ raise ValueError(
+ "The expression could not be evaluated over the given input: an "
+ "operator returned a value that the spec's `is_valid` rejected."
+ )
+ return type_spec_to_python_array(self.runtime, raw_output)
+
+
+def create_type_spec_exports(
+ runtime: _TypeSpecRuntime,
+ equations: pd.DataFrame,
+ search_output: tuple[AnyValue, AnyValue],
+ output_index: int | None,
+) -> pd.DataFrame:
+ equations = copy.deepcopy(equations)
+ _, all_out_hof = search_output
+ out_hof = all_out_hof[output_index] if output_index is not None else all_out_hof
+ expressions = []
+ callables = []
+ for _, row in equations.iterrows():
+ expression = out_hof.members[row["complexity"] - 1].tree
+ expressions.append(expression)
+ callables.append(CallableJuliaExpression(expression, runtime))
+ return pd.DataFrame(
+ {"julia_expression": expressions, "lambda_format": callables},
+ index=equations.index,
+ )
+
+
+def prepare_type_spec_fit_data(
+ model: Any,
+ X: Any,
+ y: Any,
+ Xresampled: Any,
+ weights: Any,
+ variable_names: Any,
+ complexity_of_variables: Any,
+ X_units: Any,
+ y_units: Any,
+) -> tuple[np.ndarray, np.ndarray, None, Any, np.ndarray, Any, Any, Any]:
+ if Xresampled is not None or model.denoise or model.select_k_features:
+ raise NotImplementedError(
+ "TypeSpec does not support resampling, denoising, or feature selection."
+ )
+ if isinstance(X, pd.DataFrame):
+ if variable_names is not None:
+ warnings.warn(
+ "`variable_names` has been reset to `None` as `X` is a DataFrame."
+ )
+ variable_names = X.columns.astype(str).to_numpy()
+ X = X.to_numpy(dtype=object)
+ else:
+ X = object_array_2d(X)
+ if X.ndim != 2:
+ raise ValueError("TypeSpec X must be a 2D array of logical values.")
+ if X.shape[1] == 0:
+ raise ValueError("TypeSpec X must contain at least one feature.")
+ if variable_names is not None and len(variable_names) != X.shape[1]:
+ raise ValueError("`variable_names` must provide one name per TypeSpec feature.")
+ if variable_names is not None and any(" " in name for name in variable_names):
+ variable_names = [name.replace(" ", "_") for name in variable_names]
+ warnings.warn(
+ "Spaces in variable names are not supported. "
+ "Spaces have been replaced with underscores. \n"
+ "Please use valid names instead."
+ )
+
+ if isinstance(y, (pd.Series, pd.DataFrame)):
+ y = y.to_numpy(dtype=object)
+ else:
+ y = object_array_1d(y)
+ if y.ndim == 2 and y.shape[1] == 1:
+ y = y[:, 0]
+ if y.ndim != 1:
+ raise NotImplementedError("TypeSpec currently supports one output.")
+ if X.shape[0] != y.shape[0]:
+ raise ValueError("X and y have inconsistent numbers of samples.")
+ if X.shape[0] == 0:
+ raise ValueError("X and y must contain at least one sample.")
+ if weights is not None:
+ raise NotImplementedError("TypeSpec does not currently support weights.")
+ if X_units is not None or y_units is not None:
+ raise NotImplementedError("TypeSpec does not currently support units.")
+
+ model.n_features_in_ = X.shape[1]
+ if variable_names is None:
+ variable_names = np.array([f"x{i}" for i in range(X.shape[1])])
+ model.feature_names_in_ = np.asarray(variable_names, dtype=str)
+ model.display_feature_names_in_ = model.feature_names_in_
+ model.nout_ = 1
+ model.complexity_of_variables_ = copy.deepcopy(complexity_of_variables)
+ model.X_units_ = copy.deepcopy(X_units)
+ model.y_units_ = copy.deepcopy(y_units)
+ return (
+ X,
+ y,
+ None,
+ weights,
+ model.feature_names_in_,
+ complexity_of_variables,
+ X_units,
+ y_units,
+ )
+
+
+def prepare_type_spec_prediction_data(model: Any, X: Any) -> np.ndarray:
+ if not isinstance(X, pd.DataFrame):
+ X = object_array_2d(X)
+ if X.ndim != 2:
+ raise ValueError("X must be a 2D array.")
+ if X.shape[1] != model.n_features_in_:
+ raise ValueError("X has a different number of features than during fit.")
+ X = pd.DataFrame(X)
+ if isinstance(X.columns, pd.RangeIndex):
+ if model.selection_mask_ is not None:
+ X = X[X.columns[model.selection_mask_]]
+ X.columns = model.feature_names_in_
+ columns = X.columns.astype(str)
+ if columns.str.contains(" ").any():
+ X = X.copy()
+ X.columns = columns.str.replace(" ", "_")
+ warnings.warn(
+ "Spaces in DataFrame column names are not supported. "
+ "Spaces have been replaced with underscores. \n"
+ "Please rename the columns to valid names."
+ )
+ X = X.rename(columns=str)
+ missing_features = set(model.feature_names_in_) - set(X.columns)
+ if missing_features:
+ raise ValueError(f"X is missing features: {sorted(missing_features)}")
+ return np.asarray(X.reindex(columns=model.feature_names_in_).to_numpy(dtype=object))
+
+
+def validate_type_spec_model_configuration(model: Any) -> None:
+ if model.binary_operators is not None or model.unary_operators is not None:
+ raise ValueError(
+ "TypeSpec requires `operators={...}` and does not accept "
+ "`binary_operators` or `unary_operators`."
+ )
+ validate_type_spec_configuration(
+ model.type_spec,
+ model.operators,
+ elementwise_loss=model.elementwise_loss,
+ loss_function=model.loss_function,
+ loss_function_expression=model.loss_function_expression,
+ )
+ unsupported = {
+ "turbo": model.turbo,
+ "bumper": model.bumper,
+ "autodiff_backend": model.autodiff_backend is not None,
+ "output_jax_format": model.output_jax_format,
+ "output_torch_format": model.output_torch_format,
+ "extra_sympy_mappings": model.extra_sympy_mappings is not None,
+ "extra_jax_mappings": model.extra_jax_mappings is not None,
+ "extra_torch_mappings": model.extra_torch_mappings is not None,
+ }
+ configured = [name for name, enabled in unsupported.items() if enabled]
+ if configured:
+ raise ValueError(
+ "TypeSpec does not support "
+ + ", ".join(f"`{name}`" for name in configured)
+ + "."
+ )
+ expression_spec = model.expression_spec_
+ if not expression_spec.supports_type_spec:
+ raise ValueError(f"{type(expression_spec).__name__} does not support TypeSpec.")
+ expression_spec._validate_type_spec()
diff --git a/release-please-config.json b/release-please-config.json
index 870fa18da..bb49de6be 100644
--- a/release-please-config.json
+++ b/release-please-config.json
@@ -4,8 +4,8 @@
".": {
"release-type": "python",
"include-component-in-tag": false,
- "versioning": "prerelease",
- "prerelease": true
+ "versioning": "default",
+ "prerelease": false
}
}
}
diff --git a/scripts/paper_image_sync.py b/scripts/paper_image_sync.py
index 486ea6926..b037dbf10 100644
--- a/scripts/paper_image_sync.py
+++ b/scripts/paper_image_sync.py
@@ -551,7 +551,7 @@ def main() -> None:
base=docs_default_branch,
title=f"Add paper images from PySR PR #{pr_number}",
body=(
- f"Automated upload of paper image(s) from MilesCranmer/PySR PR #{pr_number}.\n\n"
+ f"Automated upload of paper image(s) from astroautomata/PySR PR #{pr_number}.\n\n"
f"Source head: {pr.head_repo}@{pr.head_sha}\n"
f"Target dir: `{docs_images_dir}`\n"
),
diff --git a/skills/pysr/SKILL.md b/skills/pysr/SKILL.md
new file mode 100644
index 000000000..ac34153f4
--- /dev/null
+++ b/skills/pysr/SKILL.md
@@ -0,0 +1,366 @@
+---
+name: pysr
+description: Use when fitting equations to data with PySR or SymbolicRegression.jl, when a user wants an interpretable formula, symbolic model, scaling law, or empirical relation discovered from numeric data, or when debugging a PySR search that is slow, stuck, or giving poor equations.
+---
+
+# Using PySR Effectively
+
+PySR evolves expression trees through the Julia backend SymbolicRegression.jl, returning readable equations on an accuracy/complexity Pareto front. This guide distills the documentation and several hundred real user threads. Full docs: https://pysr.ai/
+
+## Quick start
+
+```python
+# pip install pysr (Julia is downloaded automatically on first import; no separate install)
+import numpy as np
+from pysr import PySRRegressor
+
+X = 2 * np.random.randn(100, 5)
+y = 2 * np.cos(X[:, 3]) + X[:, 0] ** 2 - 2
+
+model = PySRRegressor(
+ binary_operators=["+", "-", "*", "/"],
+ unary_operators=["cos"],
+ niterations=100,
+)
+model.fit(X, y)
+print(model) # Pareto front: complexity, loss, equation
+model.predict(X) # uses the auto-selected "best" equation
+model.sympy() # SymPy expression
+model.latex() # LaTeX string
+```
+
+`model.equations_` is a pandas DataFrame containing `complexity`, `loss`, `score`, `equation`, `sympy_format`, `lambda_format`. Select any front row by index in `predict`, `sympy`, `latex`, `jax`, `pytorch`; without an index, methods use the auto-selected equation. Indexing lets you inspect and export a specific Pareto candidate instead of only the default row.
+
+## Keep the process alive
+
+The first `.fit()` pays Julia startup/JIT, roughly 1-2 minutes; subsequent same-process fits start almost instantly. Repeated fresh processes repay startup and compilation for every experiment, a common agent failure.
+
+- Keep a long-lived IPython/Jupyter kernel or REPL. New `PySRRegressor` instances are cheap; Julia runtime state lives per-process.
+- Scripts work. Iterate parameters inside one process instead of repeatedly relaunching `python script.py`. Process reuse is an optimization rather than a requirement.
+- Reuse an environment containing `pysr`; fresh environments re-resolve and precompile Julia packages.
+- Expect “Compiling Julia backend...” initially. Avoid treating normal compilation as a hang.
+
+## Candidate equations with `guesses`
+
+Inject knowledge through modifiable/discardable `guesses` or mandatory template structure. Combine both using component/parameter guesses. Constructor strings require Julia syntax, input variable names, and enabled operators:
+
+```python
+model = PySRRegressor(
+ operators={2: ["+", "*"]},
+ guesses=["x0 + 0.5 * x1", "x0 * (x1 + 1.0)"],
+)
+```
+
+Single output: candidate list. Multiple outputs: lists in target-column order, `guesses=[["x0 + x1"], ["x0 * x1"]]`. Templates: dictionaries keyed by component/parameter names; `#1`, `#2`, etc. identify component arguments.
+
+Guesses enter populations throughout search. `should_optimize_constants=True` optimizes constants/template parameters before insertion. `fraction_replaced_guesses`, default `0.001`, controls the cycle-end replacement fraction.
+
+Guesses are pushed into the hall of fame before generation one. Inspect accepted seeds at their complexity in `model.equations_` or immediate `hall_of_fame.csv`; confirm insertion and optimized constants. Oversize seeds produce only Julia warning `Guess expression '...' has complexity 23 > maxsize (10)`, using default `x1, x2, ...` names regardless of supplied names. They remain unrecorded but still migrate; descendants become recordable after mutations satisfy `maxsize`. Violations of `constraints`, `nested_constraints`, or `maxdepth` disappear silently. Lower-loss candidates replace seeds at the same complexity: only one expression survives per complexity.
+
+Keep constraints compatible and set `maxsize` slightly above the largest seed; use the warning's complexity. Cost rises steeply, PySR warns above 40, and several-times-seed sizes can exhaust the budget without completing one generation despite appearing alive. Use `warmup_maxsize_by` against early bloat.
+
+Set guesses on the estimator, never `.fit()`. Continue with `model.set_params(guesses=[...], warm_start=True)`, replacing guesses while keeping data representation/search space fixed.
+
+## Recommended workflow
+
+1. **Subsample:** usually 1,000-5,000 representative rows suffice even from millions. Many features, heavy noise, or rare regimes benefit from more. PySR 2.0 automatically batches larger datasets; `batching=False` forces full-data search evaluations. Batching still reevaluates hall-of-fame candidates fully. Fewer rows proportionally accelerate search.
+2. **Minimize operators:** select domain-plausible operators. Redundant `pow`/`square`/`cube` or `-`/`neg` enlarge search. Polynomials: `["+", "-", "*"]`, omit `/` and `^`.
+3. **Use defaults:** populations, parsimony, mutation weights, and `ncycles_per_iteration` were tuned by large-scale search in 2024-2025. Avoid older forum/paper recipes.
+4. **Debug briefly:** check operators, loss, and sensible equations in a few-minute run before one long search.
+5. **Budget the final run:** evolution lacks conventional convergence and can jump families after hours. Set large `niterations`; control wall-clock with `timeout_in_seconds`, total evaluations with `max_evals`, or stopping with `early_stop_condition="stop_if(loss, complexity) = loss < 1e-6 && complexity < 10"` (Julia `(loss, complexity)` function). `max_evals` supports compute-matched comparisons; `niterations * populations * population_size` does not count evaluations. Monitor partial `outputs//hall_of_fame.csv`; stop gracefully with `q` then Enter in IPython.
+6. **Inspect the entire front:** validate candidates and report tradeoffs.
+
+## Choosing an equation
+
+`model_selection="best"` maximizes `score` within 1.5x minimum loss; `"accuracy"` minimizes loss. Score is negative log-loss slope per complexity unit; a large score marks a sharp accuracy gain for little added complexity. These heuristics control default `predict`.
+
+Print `model.equations_`; inspect sharp loss drops, held-out `model.predict(X_test, index=i)` when overfitting is plausible, and limits x -> 0/inf against domain expectations. Present 2-3 candidates with tradeoffs.
+
+## Losses and weights
+
+Default: MSE.
+
+- Outliers: `elementwise_loss="L1DistLoss()"`, median-seeking versus mean-seeking.
+- Uncertainty: `model.fit(X, y, weights=1/sigma**2)`; built-ins apply weights automatically. Custom: `elementwise_loss="myloss(x, y, w) = w * abs(x - y)^2"`.
+- Wide target range: largest values dominate MSE; use log-space loss:
+
+```python
+elementwise_loss = """function loss_fnc(prediction, target)
+ scatter_loss = abs(log((abs(prediction)+1e-20) / (abs(target)+1e-20)))
+ sign_loss = 10 * (sign(prediction) - sign(target))^2
+ return scatter_loss + sign_loss
+end"""
+```
+
+- Relative/percentage error: divide by target; prediction denominators reward infinite predictions. Guard near-zero targets.
+- LossFunctions.jl strings work: `"HuberLoss(1.0)"`, `"LPDistLoss{3}()"`.
+- Binary classification: +1/-1 targets, margin loss `"L2MarginLoss()"`; apply sigmoid yourself for probabilities.
+- Asymptotes/boundaries, including y -> 0 as x -> inf, exact boundary values, and known limits: add synthetic points with very large `weights`, usually preferable to custom objectives. Strict enforcement: estimate limits numerically inside `loss_function`, add graded penalties.
+- Keep losses deterministic for caching and non-negative; `loss_scale="linear"` permits negatives, including log-likelihoods.
+
+`elementwise_loss` receives scalar `(prediction, target)` or `(prediction, target, weight)`; never sum/broadcast. Whole-vector/tree objectives require `loss_function`.
+
+## Scaling and features
+
+- Normalization is optional: constants sample near N(0,1), mutate multiplicatively; extreme scales slow search. Prefer natural units, avoiding nuisance constants and hidden physical meaning. Rescale visibly struggling searches.
+- Roughly 10 features or fewer need no special handling.
+- Tens: increase `maxsize` and rows; search selects features reasonably well, automatic batching handles volume. Requiring all 30+ features needs `maxsize` well above 100.
+- Above ~50: engineer domain aggregates or template decompositions. Fields/images/sequences/graphs need meaningful features rather than pixel columns, or appropriately biased neural networks followed by symbolic component distillation, arXiv:2006.11287. Fallback: gradient-boosting preselection, `select_k_features=k`.
+- Omitted variables failed to justify complexity through accuracy; force inclusion through an absence-penalizing custom loss.
+
+## PDE discovery
+Regress `u_t = f(u, u_x, u_xx, ...)` from `u(x, t)`: grid points become samples, field/spatial derivatives become features, time derivative becomes target; PySR discovers products like `u*u_x` without a fixed library.
+- Spatial derivatives: spectral differentiation for clean periodic grids, otherwise Savitzky-Golay, `savgol_filter(U, 21, 3, deriv=1, delta=dx, axis=-1)`; finite differences require clean data, spectral derivatives catastrophically amplify noise.
+- Target: central snapshot differences; smooth noisy/sparse time samples along time or sample finer. Trim non-periodic stencil margins.
+- Bias low-order terms with column-ordered `complexity_of_variables=[1, 2, 3, ...]`, costing `u`, `u_x`, `u_xx` as 1, 2, 3; prune dimensional impossibilities with `X_units=["m/s", "s^-1", ...]`. Inspect the front elbow; noisy coefficients are approximate. Validate forward simulation from held-out initial conditions.
+## Template expressions
+Default to plain search; suggest `TemplateExpressionSpec` for known structure that free-form trees would rediscover or violate:
+- Outer forms: `sin(f(x1, x2)) + g(x3)`, rational functions, envelopes `x*(1-x)*g(x)`.
+- Per-category/class/object/condition coefficients through `parameters`; shared subexpressions/coupled outputs; derivative/integral relations through differential operator `D`; hard variable-placement requirements.
+```python
+from pysr import PySRRegressor, TemplateExpressionSpec
+spec = TemplateExpressionSpec(
+ combine="sin(f(x1, x2)) + g(x3)^2",
+ expressions=["f", "g"],
+ variable_names=["x1", "x2", "x3"],
+)
+model = PySRRegressor(expression_spec=spec, binary_operators=["+", "-", "*", "/"])
+model.fit(X, y)
+```
+Pass category as an X column; increment zero-based IDs for Julia's 1-based indexing:
+```python
+spec = TemplateExpressionSpec(
+ combine="p[class] * f(x1, x2) + q[class]",
+ expressions=["f"],
+ variable_names=["x1", "x2", "class"],
+ parameters={"p": 3, "q": 3},
+)
+model = PySRRegressor(
+ expression_spec=spec,
+ operators={2: ["+", "*"]},
+ guesses=[{
+ "f": "#1 + #2",
+ "p": [5.0, 10.0, 0.8],
+ "q": [0.0, 0.0, 0.0],
+ }],
+)
+```
+`parameters={"p": 3, "q": 3}` declares vector lengths; guesses initialize vectors and `f`. Declare disjoint component/parameter names and every guess key; supply every component formula. Match parameter lengths using Python lists/1D NumPy arrays; omit vectors for normal initialization. Values remain learnable, are copied without modifying caller arrays, and use model precision. Parameter guesses require SymbolicRegression.jl ≥2.3.0. `#1`/`#2` denote passed arguments, here `x1`/`x2`; with `combine="f(x, p[1])"`, `#2` denotes `p[1]`.
+`combine` accepts arbitrary Julia/multiple statements: reuse `fx = f(x); fx + fx^2`, change arguments `f(x1) - f(x2)`, differentiate `df = D(f, 1); df(x)`. Vector/multi-output problems: append extra targets to X, return per-row residuals, fit dummy y with `elementwise_loss="(p, t) -> p"` so output itself is loss.
+PySR 2.0 caveats: arbitrary Julia prevents `.sympy()`, `.latex()`, `.jax()`, `.pytorch()` export; manually reassemble component strings from `model.equations_`. Arguments print positionally as `#1, #2` because components may receive different inputs. Combine values are `ValidVector`s (`.x` data, `.valid` flag); arithmetic propagates validity, custom manipulation must unwrap/rebuild `ValidVector(raw, valid)`. Use Float32-safe `0.5f0` or explicit conversion; bare Float64 `0.5` risks type instability. Custom objectives require `loss_function_expression`. Learnable parameters use `TemplateExpressionSpec(parameters=...)`; pre-1.4 `function_symbols`/lambda-style combine API is removed in 2.0.
+
+## Custom value types: `TypeSpec`
+`TypeSpec` generates Julia structs for nodes, independently of `expression_spec`, including fixed structures. Supply fields in declaration order: `triple`, then `code`; the nested tuple remains one field. Keep X 2D/y 1D and assign object cells individually to prevent extra NumPy axes:
+```python
+import numpy as np
+from pysr import PySRRegressor, TemplateExpressionSpec, TypeSpec
+value = ((1.0, 2.0, 3.0), np.uint32(7))
+X = np.empty((1, 1), dtype=object)
+y = np.empty(1, dtype=object)
+X[0, 0] = value
+y[0] = value
+packet = TypeSpec(
+ "Packet",
+ fields={"triple": "NTuple{3,Float64}", "code": "UInt32"},
+ sample="rng -> Packet(ntuple(_ -> randn(rng), 3), rand(rng, UInt32))",
+ mutate="""function (rng, value::Packet, temperature::Float64)
+ triple = ntuple(i -> value.triple[i] + temperature * randn(rng), 3)
+ bit = UInt32(1) << rand(rng, 0:31)
+ Packet(triple, xor(value.code, bit))
+ end""",
+ scalar_constants="value -> collect(value.triple)",
+ with_scalar_constants="(value, c) -> Packet((c[1], c[2], c[3]), value.code)",
+ is_valid="value -> all(isfinite, value.triple)",
+)
+```
+- `sample(rng) -> Packet`: constant leaves; `mutate(rng, value, temperature) -> Packet`: continuous/discrete mutation; annotations are optional.
+- Supply `scalar_constants(Packet) -> Vector{Float64}` and `with_scalar_constants(Packet, Vector{Float64}) -> Packet` together; expose optimizable scalars, preserve other fields when rebuilding.
+- `is_valid(Packet) -> Bool`: reject invalid intermediates; optional `init`: `() -> Packet`; `string`: `Packet -> AbstractString`.
+- `preamble`: Julia before generated type/hooks; `definitions`: Julia immediately after type, including constructors/type-specific methods.
+- `loss_type`: concrete `AbstractFloat` for full objectives; unavailable for elementwise losses, whose return type is inferred.
+Configure `PySRRegressor(type_spec=packet, operators=operators, elementwise_loss=loss)`. In arity-keyed `operators`, unary `negate_packet(x::Packet)::Packet = Packet(ntuple(i -> -x.triple[i], 3), x.code)` preserves code; binary `combine_packets(x::Packet, y::Packet)::Packet = Packet(ntuple(i -> x.triple[i] + y.triple[i], 3), xor(x.code, y.code))` sums triples/XORs codes. Include binary `choose_parameter(a::Packet, b::Packet) = a` and `choose_parameter(a::Packet, b::ValidVector) = ValidVector(map(_ -> a, b.x), b.valid)`. Define `loss` as `packet_loss(p::Packet, y::Packet)::Float64 = sum((p.triple[i] - y.triple[i])^2 for i in 1:3) + Float64(count_ones(xor(p.code, y.code)))`; annotation keeps return type concrete.
+Typed constructors work in `ordinary_guess = ["combine_packets(x0, Packet((1.0, 2.0, 3.0), UInt32(7)))"]` and `template_component_guess = [{"f": "combine_packets(#1, Packet((1.0, 2.0, 3.0), UInt32(7)))"}]`. Parsing uses the generated configuration module with Packet/operators defined; no parser closure/import needed.
+Template parameters share the data's `TypeSpec`. Define `packet_template = TemplateExpressionSpec(combine="choose_parameter(p[1], f(x))", expressions=["f"], variable_names=["x"], parameters={"p": 1})`. Preserve one logical Packet with `packet_guess_from_list = [{"f": "#1", "p": [value]}]`, or assign `packet_parameter_array = np.empty(1, dtype=object)` and `packet_parameter_array[0] = value`, then `packet_guess_from_array = [{"f": "#1", "p": packet_parameter_array}]`. Configure the estimator above with additional `expression_spec=packet_template`, `guesses=packet_guess_from_array`.
+Verify field count/order, conversion to declared Julia types, type-stable Packet-returning operators, and one real scalar per loss pair; reserve `Inf` for invalid evaluations.
+
+## Custom operators
+
+Supply Julia strings and SymPy mappings using SymPy functions, never numpy/scipy; incorrect mappings break export/prediction:
+
+```python
+model = PySRRegressor(
+ binary_operators=["+", "*"],
+ unary_operators=["inv(x) = 1/x", "gauss(x) = exp(-x^2)"],
+ extra_sympy_mappings={
+ "inv": lambda x: 1/x,
+ "gauss": lambda x: sympy.exp(-x**2),
+ },
+)
+```
+
+- Accept all reals, including probes beyond data. Guard domains: `my_sqrt(x) = x >= 0 ? sqrt(x) : convert(typeof(x), NaN)`; any data-point NaN discards candidates with infinite loss. Prefer protected built-ins `sqrt`, `log`, `acosh`.
+- Preserve types: `T(2.5)` with `where {T}`, or default-Float32 `2.5f0`; bare `2.5`/`0` introduce Float64/Int64 and break Float32 pipelines.
+- One/two scalar inputs, scalar output; ≥3 arguments require arity-keyed `operators` or templates.
+- Julia packages: `from pysr import jl; jl.seval("import Pkg; Pkg.add(\"SpecialFunctions\")"); jl.seval("using SpecialFunctions")`; safely wrap any importable function.
+- No closed SymPy form: map to `class myop(sympy.Function): pass`; symbolic export works, `predict` requires numeric mapping. Call `model.refresh()` after mapping changes.
+
+## Custom objectives: `loss_function`
+
+Use full objectives for vectors, trees, derivatives, auxiliary data:
+
+```python
+objective = """
+function my_objective(tree, dataset::Dataset{T,L}, options) where {T,L}
+ prediction, completed = eval_tree_array(tree, dataset.X, options)
+ !completed && return L(Inf)
+ residuals = prediction .- dataset.y
+ return sum(abs2, residuals) / dataset.n
+end
+"""
+model = PySRRegressor(loss_function=objective, binary_operators=["+", "*", "-"])
+```
+
+- Check `completed` before accessing predictions; failure leaves garbage. Return `L(Inf)` for numerical invalidity; use graded finite structural penalties such as `L(1e6 * n_violations)` so intermediate mutations survive. All-or-nothing structural rejection can make targets unreachable.
+- `dataset.X`: features × samples, transposed from Python; `dataset.y`: vector; `dataset.n`: sample count. Three-argument objectives receive active batches; `(tree, full_dataset, options, idx)` exposes full data/selected indices.
+- Millions of calls: type-stable, vectorized, no printing/Python callbacks (GIL serialization). Use `const` globals or string interpolation, Distances.jl/standard packages; diagnose with `@code_warntype`/BenchmarkTools in separate Julia. Avoid internal-name collision `eval_loss`.
+- Append auxiliary targets/derivatives/group IDs to X and slice, or interpolate literal arrays. Appended columns remain searchable unless templates exclude them or penalties discourage use.
+- Candidate derivatives: `eval_diff_tree_array` for one feature, `eval_grad_tree_array` for all; support monotonicity/physics-informed losses. Templates use `D(f, i)`.
+- Reinterpreted trees, including split subtrees/recursion, print raw equations requiring manual decoding; symbolic manipulation generally breaks `.sympy()`/`.predict()`, requiring manual extraction.
+
+### Structural tree constraints
+
+Before writing a traversal, check whether the rule is already an option. Operator placement rules are `nested_constraints`: "no sin inside cos and no cos inside sin" is `nested_constraints={"sin": {"cos": 0}, "cos": {"sin": 0}}`; "exponent must be a bare constant or variable" is `constraints={"^": (-1, 1)}` (keys are the names used in the operator lists, so `"pow"` if that is the name given). Both reject a violating mutation before it is evaluated, which is cheaper than a loss penalty and never lets a violator into the hall of fame. Use traversal only for what they cannot express: required features, constant value ranges, rules that depend on position or on several operators at once.
+
+Implement those by traversal. Supported: `any(f, tree)`, `all(f, tree)`, `count(f, tree)`, `sum(f, tree)`, `foreach(f, tree)`, `collect(tree)`, depth-first `for node in tree`. Prefer functional traversal: loops allocate stacks; `count`/`sum`/`any` traverse directly.
+
+- `node.degree`: leaf 0/unary 1/binary 2; `node.l`, `node.r`: traversable children.
+- Leaves: `node.constant` distinguishes constants/variables; `node.val`: constant value; `node.feature`: 1-based feature index.
+- `node.op`: 1-based supplied operator-list index, per arity.
+
+Allow `^` only with a lone constant exponent in [0,1]:
+
+```python
+objective = """
+function constrained_loss(tree, dataset::Dataset{T,L}, options) where {T,L}
+ idx_pow = 3 # position of ^ in binary_operators below (1-indexed)
+ n_bad = count(tree) do node
+ node.degree == 2 && node.op == idx_pow &&
+ any(c -> !(c.degree == 0 && c.constant && 0 <= c.val <= 1), node.r)
+ end
+ n_bad > 0 && return L(10_000 * n_bad)
+ prediction, valid = eval_tree_array(tree, dataset.X, options)
+ !valid && return L(Inf)
+ return sum(i -> abs2(prediction[i] - dataset.y[i]), eachindex(prediction)) / dataset.n
+end
+"""
+model = PySRRegressor(binary_operators=["+", "*", "^"], loss_function=objective)
+```
+
+Count violations before evaluation; multiply by a large finite penalty to guide evolution. Require feature 2: `any(n -> n.degree == 0 && !n.constant && n.feature == 2, tree) || return L(big)`. Templates: obtain component trees using `get_tree(ex)` inside `loss_function_expression`. `tree_mapreduce`, `NodeSampler`, construction: https://ai.damtp.cam.ac.uk/dynamicexpressions/stable/examples/base_operations/
+
+## Constraints and complexity
+
+- `constraints={"pow": (9, 1)}` caps argument complexity: base 9, lone constant/variable exponent; `-1` unlimited. Strongly constrain `^`; unrestricted exponentiation searches poorly.
+- `nested_constraints={"sin": {"sin": 0, "cos": 0}, "cos": {"sin": 0, "cos": 0}}` prohibits nested trig; values cap inner occurrences within outer operators.
+- `complexity_of_operators={"exp": 3}`, `complexity_of_constants=2` discourage operators/constants; `complexity_of_variables` accepts global/per-feature costs.
+- Constraints affect every intermediate; excessive tightness silently excludes targets. Allow slack: final size 30, `maxsize=35`.
+- `maxsize` counts operator/constant/variable nodes or `complexity_mapping` units; seven-feature linear models cost ~29, default 30 excludes many-term formulas.
+- `warmup_maxsize_by=0.5` ramps maxsize over that run fraction, helping avoid early complex traps.
+- `parsimony`/`adaptive_parsimony_scaling`: defaults well-tuned in 1.5+.
+- Known pi/G: recognize fitted floats afterwards, or supply exact constant-valued features.
+- Integers: constant features plus increased `complexity_of_constants`, or custom-operator rounding; continuous optimization does not guarantee integers.
+
+### `complexity_mapping`
+
+Replace fixed operator prices/node counts with a Julia expression-to-`Int` function for emitted characters, evaluation cost, tokens, or depth. Pricing applies to front, `maxsize`, parsimony, `constraints`.
+
+```python
+complexity_mapping = """
+function source_chars(ex)
+ tree_mapreduce(
+ leaf -> leaf.constant ? length(string(leaf.val)) : 2, # "x0" is 2 chars
+ branch -> branch.degree == 1 ? 6 : 5, # "sin(" + ")" / " + " with parens
+ (parent, children...) -> parent + sum(children),
+ get_tree(ex),
+ )
+end
+"""
+model = PySRRegressor(
+ binary_operators=["+", "-", "*"],
+ complexity_mapping=complexity_mapping,
+ maxsize=280,
+)
+```
+
+This mapping tracked every test-row `len(equation)` within offset 2 for outer parentheses; rank by actual cost and calibrate against emitted candidates.
+
+- Input: `AbstractExpression`; `get_tree(ex)` exposes `degree`, `constant`, `val`, `feature`, `op`. `tree_mapreduce`/`get_tree` are in scope.
+- Avoid `length(string_tree(ex))`: internal rendering uses `binary_operator[3](...)` and `x1, x2` placeholders, unrelated to exported length.
+- Set real-budget `maxsize`; default 30 with character pricing leaves `equations_` empty when nothing fits. Hall of fame allocates one slot/unit through maxsize; filter long fine-grained fronts.
+- Every candidate incurs computation: cheap, type-stable, no Python callbacks.
+- `guesses` use identical pricing/acceptance rules.
+
+## Dimensional constraints
+
+Check physical units during search:
+
+```python
+model.fit(X, y, X_units=["Constants.M_sun", "kg", "m"], y_units="kg * m / s^2")
+```
+
+- DynamicQuantities.jl notation; `"1"` explicitly dimensionless.
+- Soft `dimensional_constraint_penalty`, default 1000; avoid 1e9, preserve graded paths through imperfect intermediates.
+- Fitted constants have wildcard `[⋅]`/`[?]` units; a lone constant can absorb any dimensions. Forbid with `dimensionless_constants_only=True`.
+- SI radians are dimensionless; units cannot restrict variables to trig.
+
+## Parallelism
+
+- Default `parallelism="multithreading"` suits laptops/single nodes. Thread count is fixed at Julia startup: set `PYTHON_JULIACALL_THREADS=` before importing pysr; `JULIA_NUM_THREADS` is inappropriate under juliacall.
+- Keep `populations` ~2-3x threads/cores for worker utilization; default `populations=31` covers typical machines.
+- Saturated coordinating thread: increase `ncycles_per_iteration`, reducing communication.
+- Long runs: `parallelism="multiprocessing"`, `procs=n` costs substantially more startup/fit, potentially faster steady-state. Multiple nodes: `cluster_manager="slurm"` or native SlurmClusterManager.jl, https://julia.pysr.ai/dev/slurm/. Launch once on one node, let it spawn workers, avoid `srun`; put worker packages in `worker_imports`.
+- Reproducibility: `deterministic=True, random_state=, parallelism="serial"`; parallel seeds remain nondeterministic. CPU differences can persist serially; reduce with `precision=64`.
+
+## Saving, resuming, exporting
+
+- Non-temporary fits write continuously readable `outputs//hall_of_fame.csv` and `checkpoint.pkl`. Reload `PySRRegressor.from_file(run_directory=...)` with identical PySR version; pickles are version-locked. Preserve CSV/construction code durably; resupply custom `extra_sympy_mappings`.
+- `warm_start=True` resumes previous populations on next same-process `.fit()`. Fix operators, `maxsize`, `expression_spec`, precision, feature count/order; loss/weights may change for staged objectives.
+- Export `model.sympy(i)`, `model.latex(i)`, `model.latex_table()`, differentiable `model.jax(i)` returning `{'callable', 'parameters'}`, trainable `model.pytorch(i)`. Supply custom `extra_jax_mappings`/`extra_torch_mappings`; fine-tune selected equations' constants by full-data gradient descent.
+- Julia losses can differ from NumPy recomputation through default 32-bit versus 64-bit precision. Use `precision=64` when important or values exceed ~1e19/fall below ~1e-19, risking Float32 overflow.
+
+## Troubleshooting
+
+| Symptom | Cause / fix |
+|---|---|
+| First import/fit takes minutes | Initial Julia download; per-process JIT. Retain process. |
+| Compilation hangs | Apple Silicon: replace Rosetta x86 Python with native arm64. Notebook/embedded stdin wedge: `input_stream="devnull"`. |
+| Startup `lock.pid` | Killed-process juliapkg lock: verify no installation, delete lock, import once freshly. Shared filesystems: pre-initialize before parallel jobs. |
+| Jupyter `UnicodeDecodeError` | Old PythonCall bug: upgrade or set `PYTHON_JULIACALL_AUTOLOAD_IPYTHON_EXTENSION=no` before import. |
+| Cannot interrupt notebook | Known limitation; IPython/terminal `q`+Enter only. Set `timeout_in_seconds` or use IPython. |
+| Nonsensical equations | Check scale/loss (log-space), missing/redundant operators, insufficient `maxsize`, excluding constraints, in order before more compute. |
+| Tiny equations | Insufficient `maxsize` or MSE dominated by largest points. |
+| `DomainError` | Guard custom operators with typed NaN. |
+| Prediction/SymPy fails | Correct `extra_sympy_mappings`; use SymPy functions. |
+| Division near zero | Any-NaN-anywhere rejects candidates; survivors remain finite only on supplied data, potentially diverging elsewhere. |
+| Early-stop `ProcessExitedException` | Harmless multiprocessing worker-teardown noise. |
+| Growing memory | Mostly fixed in recent Julia; upgrade or set multiprocessing `heap_size_hint_in_bytes`. |
+| Different results | Stochastic evolution; apply serial determinism recipe. |
+
+## Dropping to Julia
+
+Use native SymbolicRegression.jl, `SRRegressor` via MLJ, for Julia pipelines/deep customization: expression types, mutation operators, per-component constraints. `from pysr import jl` exposes the runtime; `jl.seval(...)` executes arbitrary Julia, installed packages support operators/losses. Small readable backend extension points: `src/Options.jl`, `src/CheckConstraints.jl`; connect dev checkouts through `pysr/juliapkg.json`.
+
+## Version notes
+
+Target: PySR 2.0.0, SymbolicRegression.jl 2.x. Relative to 1.5.x: n-ary `operators={1: ["sin"], 2: ["+", "*"], 3: ["clamp"]}`, `guesses`, `TypeSpec`, automatic batching, autodiff plugins. Defaults: `batching="auto"`, `batch_size=None`, `annealing=True`, `crossover_probability=0.2`, new mutation mix; use 2.0 defaults rather than 1.x tuning.
+
+| Deprecated | Current |
+|---|---|
+| `pysr.install()`, `python -m pysr install`, PyCall/`from julia import Main` | Automatic Julia installation on import |
+| `pysr(...)`, `best()`, `get_hof()` | `PySRRegressor`, `.fit()`, `model.equations_`, `get_best()` |
+| `multithreading=True/False`, `procs=0` | `parallelism="multithreading"/"multiprocessing"/"serial"` |
+| `full_objective=` | `loss_function`; templates: `loss_function_expression` |
+| `loss=` | `elementwise_loss` |
+| `equation_file=` | `output_directory=` + `run_id=` |
+| `npop`, `ncyclesperiteration`, camelCase args | `population_size`, `ncycles_per_iteration`, snake_case |
+| Positional `TemplateExpressionSpec`, `function_symbols=...` | Explicit `combine=`, `expressions=`, `variable_names=` |
+
+
+`Letter` stores a code in `0:25`. Its `string` hook displays constants as letters, such as `V` for code `21`.
+
+```python
+from pysr import PySRRegressor, TypeSpec
+
+spec = TypeSpec(
+ "Letter",
+ fields={"code": "Int"},
+ sample="rng -> Letter(rand(rng, 0:25))",
+ mutate="(rng, value, temperature) -> Letter(mod(value.code + rand(rng, (-1, 1)), 26))",
+ string="value -> string(Char('A' + value.code))",
+)
+```
+
+
+The three operators implement addition, subtraction, and multiplication modulo $26$. Every result remains a `Letter`, and the modulus appears directly in each operator. That built-in alphabet size makes this the easier cipher search.
+
+Use absolute code distance as `elementwise_loss` to reward partial improvements. Zero loss means every ciphertext/plaintext pair matches:
+
+```python
+from pysr import PySRRegressor
+
+model = PySRRegressor(
+ type_spec=spec,
+ operators={
+ 2: [
+ "shift(a::Letter, b::Letter) = Letter(mod(a.code + b.code, 26))",
+ "unshift(a::Letter, b::Letter) = Letter(mod(a.code - b.code, 26))",
+ "mix(a::Letter, b::Letter) = Letter(mod(a.code * b.code, 26))",
+ ]
+ },
+ elementwise_loss="letter_loss(prediction::Letter, target::Letter)::Float64 = abs(prediction.code - target.code)",
+ niterations=20,
+ random_state=0,
+)
+model.fit(X, y)
+print(model.equations_[["complexity", "loss", "equation"]].to_string(index=False))
+```
+
+
+A `loss_function` is useful when scoring needs the dataset as a whole, as in the rollout in [Swinging up a cart-pole](/examples/objectives#swinging-up-a-cart-pole-with-a-rollout-objective) or the period certificate in [Inventing a pseudorandom generator](/examples/objectives#inventing-a-pseudorandom-generator-with-no-target).
+
+One recovered expression is `mix(V, shift(S, x0))`: `V` has code $21$ and `S` has code $18$, giving $21(x_0 + 18) \bmod 26$. Since $21$ is the multiplicative inverse of $5$ modulo $26$, this expression decodes the ciphertext.
+
+### Learning the modulus too
+
+The harder variant removes the alphabet size from the operators. It uses values in `0:51` with integer addition, multiplication, and modulo, so the search must discover the constant $26$ along with the decryption law. This removes a strong structural hint and enlarges the constant search.
+
+One recovered decryption law is:
+
+```
+modulo(mul(add(18, x0), 47), 26)
+```
+
+Run `examples/affine_cipher.py` for the first search or `examples/affine_freemod.py` for the full type, operator, and loss setup of this variant. Allow several minutes for the latter.
+
+## Rediscovering Conway's Game of Life
+
+
+
+Conway's Game of Life can be written as a rule about a cell and its live-neighbour count. This example withholds that rule and supplies only its transition table, asking PySR to reconstruct the logical expression.
+
+The table has $18$ distinct states: two current cell states, each paired with every neighbour count from $0$ through $8$. Repeating the table $8$ times produces $144$ rows for scoring without introducing new cases. A zero-one loss therefore charges $1/18$ for a rule that misses one distinct state. The arrays use `dtype=object` so each value can be wrapped in the custom Julia type.
+
+Data generation code
+ +```python +import numpy as np + +ALPHABET = "ABCDEFGHIJKLMNOPQRSTUVWXYZ" +PLAIN = "THELAWSOFPHYSICSAREWRITTENINTHELANGUAGEOFMATHEMATICS" +CIPHER = "".join(ALPHABET[(5 * ALPHABET.index(c) + 8) % 26] for c in PLAIN) + +X = np.array([[ALPHABET.index(c)] for c in CIPHER], dtype=object) +y = np.array([ALPHABET.index(c) for c in PLAIN], dtype=object) +``` + +
+
+
+The `Cell` type wraps a single `Int`, which is enough for the current state, the neighbour count, and the output code. Its hooks propose small counts, move existing constants locally modulo $9$, and print the integer payload.
+
+```python
+from pysr import PySRRegressor, TypeSpec
+
+spec = TypeSpec(
+ "Cell",
+ fields={"v": "Int"},
+ sample="rng -> Cell(rand(rng, 0:8))",
+ mutate="(rng, value, temperature) -> Cell(mod(value.v + rand(rng, (-1, 1)), 9))",
+ string="value -> string(value.v)",
+)
+```
+
+
+The operator vocabulary contains `not`, `and`, `or`, and the equality test `eq`. Arithmetic and threshold comparisons are absent, so a neighbour threshold can enter an expression only as a constant in a call such as `eq(n, 3)`. The loss is zero-one equality on each table row. It records exact agreement, which suits a discrete truth table where the numerical distance between cell codes carries no meaning.
+
+```python
+from pysr import PySRRegressor
+
+model = PySRRegressor(
+ type_spec=spec,
+ operators={
+ 1: ["not(a::Cell) = Cell(a.v == 0 ? 1 : 0)"],
+ 2: [
+ "and(a::Cell, b::Cell) = Cell((a.v != 0 && b.v != 0) ? 1 : 0)",
+ "or(a::Cell, b::Cell) = Cell((a.v != 0 || b.v != 0) ? 1 : 0)",
+ "eq(a::Cell, b::Cell) = Cell(a.v == b.v ? 1 : 0)",
+ ],
+ },
+ elementwise_loss="cell_loss(prediction::Cell, target::Cell)::Float64 = prediction.v == target.v ? 0.0 : 1.0",
+ niterations=160,
+ random_state=0,
+)
+model.fit(X, y, variable_names=["alive", "n"])
+print(model.equations_[["complexity", "loss", "equation"]].to_string(index=False))
+```
+
+
+A recovered rule is:
+
+```
+or(eq(n, 3), and(eq(n, 2), alive))
+```
+
+The first branch allows birth or survival at three neighbours. The second adds survival of a live cell at two neighbours. It agrees with all 18 possible input states, so the transition table fully checks this local rule.
+
+Run `examples/game_of_life.py` to reproduce the Game of Life search.
+
+## Searching over machine words
+
+
+
+This search targets the bit interleave used by Morton, or Z-order, indexing. It gives every expression a `Word` value and restricts the vocabulary to machine instructions: xor, and, or, not, left shift, and right shift. Shift distances are ordinary `Word` constants, so the search must find them as part of the expression.
+
+The reference interleave is evaluated on all $256$ pairs of $4$-bit inputs.
+
+Data generation code
+ +```python +import numpy as np + +TABLE = [(a, n, int(n == 3 or (a == 1 and n == 2))) for a in (0, 1) for n in range(9)] +REPLICAS = 8 +X = np.array([[a, n] for a, n, _ in TABLE] * REPLICAS, dtype=object) +y = np.array([t for _, _, t in TABLE] * REPLICAS, dtype=object) +``` + +Data generation code
+ +```python +import numpy as np + +from pysr import PySRRegressor, TypeSpec + +WIDTH = 4 +WIDER = 5 + + +def morton(x, y, width=WIDTH): + z = 0 + for i in range(width): + z |= ((x >> i) & 1) << (2 * i) + z |= ((y >> i) & 1) << (2 * i + 1) + return z + + +PAIRS = [(a, b) for b in range(1 << WIDTH) for a in range(1 << WIDTH)] +X = np.array(PAIRS, dtype=object) +y = np.array([morton(a, b) for a, b in PAIRS], dtype=object) +``` + +The chamfer loss and the raster
+ +```julia +# Symmetric chamfer distance in pixels, divided by the grid width. Zero means +# the two drawings lit exactly the same pixels. +function _chamfer(pk::Vector{Int8}, pv::Vector{Float64}, + tk::Vector{Int8}, tv::Vector{Float64}) + a = _raster(pk, pv) + a === nothing && return _FAIL + b, db = _target_grids(tk, tv) + any(b) || return _FAIL + da = _edt(a) + ab = _mean_to(a, db); ba = _mean_to(b, da) + (ab === nothing || ba === nothing) && return _FAIL + v = 0.5 * (ab + ba) / _GRID + return isfinite(v) ? v : _FAIL +end +``` + +
+
+
+`FORCE` declares the three fields and the hooks that control sampling, constant
+optimization, string output, and failed evaluation. `sample` draws a value for each
+field. The paired `scalar_constants` and `with_scalar_constants` hooks expose and restore
+all three components, so every `Force` constant is a vector. `init_invalid` supplies a
+`Force` with `NaN` in every field. No explicit `is_valid` is given, so the
+default finite-constant check rejects that value:
+
+```python
+from pysr import TypeSpec
+
+FORCE = TypeSpec(
+ "Force",
+ fields={"x": "Float64", "y": "Float64", "z": "Float64"},
+ sample="rng -> Force(randn(rng, 3)...)",
+ definitions="""
+ Base.sin(a::Force)::Force = Force(sin(a.x), sin(a.y), sin(a.z))
+ Base.cos(a::Force)::Force = Force(cos(a.x), cos(a.y), cos(a.z))
+ function Base.sqrt(a::Force)::Force
+ f(v) = v < 0 ? NaN : sqrt(v)
+ return Force(f(a.x), f(a.y), f(a.z))
+ end
+ Base.exp(a::Force)::Force = Force(exp(a.x), exp(a.y), exp(a.z))
+ Base.:+(a::Force, b::Force)::Force = Force(a.x + b.x, a.y + b.y, a.z + b.z)
+ Base.:-(a::Force, b::Force)::Force = Force(a.x - b.x, a.y - b.y, a.z - b.z)
+ Base.:*(a::Force, b::Force)::Force = Force(a.x * b.x, a.y * b.y, a.z * b.z)
+ Base.:/(a::Force, b::Force)::Force = Force(a.x / b.x, a.y / b.y, a.z / b.z)
+ """,
+ scalar_constants="value -> [value.x, value.y, value.z]",
+ with_scalar_constants="(value, c) -> Force(c[1], c[2], c[3])",
+ string="value -> sprint(show, (value.x, value.y, value.z); context = :compact => true)",
+ init_invalid="() -> Force(NaN, NaN, NaN)",
+)
+```
+
+`definitions` runs after the `Force` type is generated, so its Julia methods can return
+`Force` values. The guarded `sqrt` returns `NaN` for negative fields instead of throwing.
+Register the methods by arity and define a loss over all three force components:
+
+```python
+OPERATORS = {
+ 1: [
+ "Base.sin",
+ "Base.cos",
+ "Base.sqrt",
+ "Base.exp",
+ ],
+ 2: [
+ "Base.:+",
+ "Base.:-",
+ "Base.:*",
+ "Base.:/",
+ ],
+}
+
+FORCE_LOSS = "force_loss(a::Force, b::Force)::Float64 = (a.x - b.x)^2 + (a.y - b.y)^2 + (a.z - b.z)^2"
+```
+
+Inside a template `combine`, PySR's backend wraps input columns and named-hole results in
+`ValidVector` values. The wrapper exposes `.x` for raw values and `.valid` for evaluation
+status; `ValidVector(raw, valid)` constructs a wrapper for a new result. This type is supplied
+by the template runtime rather than defined in this example.
+
+`PHYSICS` checks the hole results for validity, then assembles one `Force` per row
+using the fixed cross product and drag law:
+
+```python
+from pysr import TemplateExpressionSpec
+
+PHYSICS = r"""
+begin
+ _B_x = B_x(t)
+ _B_y = B_y(t)
+ _B_z = B_z(t)
+ _F_d_scale = F_d_scale(T)
+ if !(_B_x.valid && _B_y.valid && _B_z.valid && _F_d_scale.valid)
+ return ValidVector(_B_x.x, false)
+ end
+ F = map(_B_x.x, _B_y.x, _B_z.x, _F_d_scale.x,
+ v_x.x, v_y.x, v_z.x) do bx, by, bz, fd, ux, uy, uz
+ b1, b2, b3 = bx.x, by.x, bz.x
+ u1, u2, u3 = ux.x, uy.x, uz.x
+ s = fd.x
+ Force(u2 * b3 - u3 * b2 + s * u1,
+ u3 * b1 - u1 * b3 + s * u2,
+ u1 * b2 - u2 * b1 + s * u3)
+ end
+ ValidVector(F, true)
+end
+"""
+
+STRUCTURE = TemplateExpressionSpec(
+ combine=PHYSICS,
+ expressions=["B_x", "B_y", "B_z", "F_d_scale"],
+ variable_names=variable_names,
+)
+```
+
+The field holes receive only `t`; `F_d_scale` receives only `T`. Only the `.x` fields
+of their `Force` values enter the numerical residual. The `.y` and `.z` fields can still
+affect whether intermediate values remain valid.
+
+```python
+from pysr import PySRRegressor
+
+model = PySRRegressor(
+ type_spec=FORCE,
+ expression_spec=STRUCTURE,
+ operators=OPERATORS,
+ elementwise_loss=FORCE_LOSS,
+ niterations=100,
+ verbosity=0,
+ random_state=0,
+)
+
+model.fit(X, y, variable_names=variable_names)
+print(model.equations_[["complexity", "loss", "equation"]].to_string(index=False))
+```
+
+An example fitted expression is:
+
+ B_x = sin(#1 / (0.159155, 6.43746, -0.55788))
+ B_y = cos(#1 / (0.159155, 6.43746, -0.55788))
+ B_z = exp(#1 * (-0.1, -0.334652, -1.00446))
+ F_d_scale = -0.000172745, 0.262294, 0.471951
+
+`#1` is the first argument passed to each hole: `t` for the field components and
+`T` for the drag scale. Constants print three entries because they are full `Force`
+values. The field expressions match the sinusoidal and exponential forms used to
+generate the data; the constant drag scale only approximates the temperature-dependent law.
+
+The script validates force predictions on independently generated experiments with
+`held_out_error(model)`. Accurate force predictions alone do not establish exact
+recovery of every coefficient function. Expect an expensive search: a run took about
+1.5 hours in the recorded full-vector configuration.
diff --git a/docs/src/examples/getting-started.md b/docs/src/examples/getting-started.md
new file mode 100644
index 000000000..5b380d4cd
--- /dev/null
+++ b/docs/src/examples/getting-started.md
@@ -0,0 +1,153 @@
+# Getting started
+
+## Simple search
+
+This target combines a cosine term, a square term, and a constant.
+
+The model searches with `+`, `-`, `*`, and `/`. To include the cosine used in the target, add `unary_operators=["cos"]`.
+
+```python
+import numpy as np
+
+from pysr import PySRRegressor
+
+X = 2 * np.random.randn(100, 5)
+y = 2 * np.cos(X[:, 3]) + X[:, 0] ** 2 - 2
+model = PySRRegressor(binary_operators=["+", "-", "*", "/"])
+model.fit(X, y)
+print(model)
+```
+
+## Custom operator
+
+A custom operator has a definition for the symbolic search and a callable for SymPy export and for evaluating equations. Here `inv(x) = 1/x` defines the Julia-side reciprocal operator, while `extra_sympy_mappings` gives that name its Python implementation.
+
+```python
+import numpy as np
+
+from pysr import PySRRegressor
+
+X = 2 * np.random.randn(100, 5)
+y = 1 / X[:, 0]
+model = PySRRegressor(
+ binary_operators=["+", "*"],
+ unary_operators=["inv(x) = 1/x"],
+ extra_sympy_mappings={"inv": lambda x: 1/x},
+)
+model.fit(X, y)
+print(model)
+```
+
+The generic `operators` parameter replaces the separate `unary_operators` and `binary_operators`. Its integer keys specify arity, so a dictionary can describe unary, binary, and higher-arity operators. See the [Operators of any arity](/examples/search-behaviour#operators-of-any-arity) section for this form.
+
+## Multiple outputs
+
+A two-dimensional target makes the model multi-output. Here each column of `y` is the reciprocal of a different input feature, so the fitted model supplies one selected equation for each output.
+
+```python
+import numpy as np
+
+from pysr import PySRRegressor
+
+X = 2 * np.random.randn(100, 5)
+y = 1 / X[:, [0, 1, 2]]
+model = PySRRegressor(
+ binary_operators=["+", "*"],
+ unary_operators=["inv(x) = 1/x"],
+ extra_sympy_mappings={"inv": lambda x: 1/x},
+)
+model.fit(X, y)
+```
+
+## Plotting an expression
+
+After fitting the multi-output model, `model.latex()[0]` returns the LaTeX form of the selected equation for output 0.
+
+```python
+model.latex()[0]
+```
+
+Use `model.latex()[1]` for output 1.
+
+The plot below compares the target and prediction for output 0.
+
+```python
+from matplotlib import pyplot as plt
+plt.scatter(y[:, 0], model.predict(X)[:, 0])
+plt.xlabel('Truth')
+plt.ylabel('Prediction')
+plt.show()
+```
+
+
+
+The `index` argument chooses a candidate equation from `model.equations_`. Pass it to `predict`, `sympy`, or `latex` to inspect or evaluate another candidate. For multiple outputs, provide one equation index per output in output order, then select the desired output column from the predictions.
+
+## Feature selection
+
+Feature selection reduces the candidate input columns before symbolic regression. Set `select_k_features=5` to use a random-forest preprocessor that retains up to five candidates. For structured high-dimensional data, the paper [2006.11287](https://arxiv.org/abs/2006.11287) discusses breaking the problem into smaller pieces before applying PySR.
+
+The target in this example uses features 3 and 19.
+
+```python
+import numpy as np
+
+X = np.random.randn(300, 30)
+y = X[:, 3]**2 - X[:, 19]**2 + 1.5
+```
+
+Configure feature selection in the model:
+
+```python
+from pysr import PySRRegressor
+
+model = PySRRegressor(
+ binary_operators=["+", "-", "*", "/"],
+ unary_operators=["exp"],
+ select_k_features=5,
+)
+```
+
+Fit the model to select the input columns before the search:
+
+```python
+model.fit(X, y)
+```
+
+PySR prints the selected feature names before starting the search, for example:
+
+```text
+Using features ['x3', 'x5', 'x7', 'x19', 'x21']
+```
+
+The names use zero-based column indices. This selection includes both target features, `x3` and `x19`, plus three other candidates. The symbolic search uses only these retained columns.
+
+## Denoising
+
+When per-observation weights are available, pass `weights` to `fit` and define `elementwise_loss` with the matching signature. Without weights, the function accepts `(prediction, target)`. With weights, it accepts `(prediction, target, weight)`. The weighted example uses `elementwise_loss="myloss(x, y, w) = w * (x - y)^2"`. Pass a full custom objective as `loss_function`.
+
+The `denoise=True` option preprocesses the targets with a Gaussian process before symbolic regression. Its predictions become the targets searched by PySR. Because this example passes no `Xresampled`, the predictions are made at the original rows in `X`.
+
+The target combines an exponential term, two linear terms, and random noise:
+
+```python
+import numpy as np
+
+X = np.random.randn(100, 5)
+noise = np.random.randn(100) * 0.1
+y = np.exp(X[:, 0]) + X[:, 1] + X[:, 2] + noise
+```
+
+Create and fit the denoising model:
+
+```python
+from pysr import PySRRegressor
+
+model = PySRRegressor(
+ binary_operators=["+", "-", "*", "/"],
+ unary_operators=["exp"],
+ denoise=True,
+)
+model.fit(X, y)
+print(model)
+```
diff --git a/docs/src/examples/instrumentation.md b/docs/src/examples/instrumentation.md
new file mode 100644
index 000000000..035440646
--- /dev/null
+++ b/docs/src/examples/instrumentation.md
@@ -0,0 +1,227 @@
+# Instrumentation and workflow
+
+## Using TensorBoard for logging
+
+Pass `TensorBoardLoggerSpec` through `logger_spec` to write search progress to TensorBoard. `TensorBoardLoggerSpec` records the model's search hyperparameters and writes scalar summaries for the run. The example uses `min_loss`, which tracks the current best loss, and `pareto_volume`, which summarizes the Pareto front.
+
+```python
+import numpy as np
+
+from pysr import PySRRegressor, TensorBoardLoggerSpec
+
+rstate = np.random.RandomState(42)
+
+# Uniform dist between -3 and 3:
+X = rstate.uniform(-3, 3, (1000, 2))
+y = np.exp(X[:, 0]) + X[:, 1]
+
+# Create a logger that writes to "logs/run*":
+logger_spec = TensorBoardLoggerSpec(
+ log_dir="logs/run",
+ log_interval=10, # Log every 10 iterations
+)
+
+model = PySRRegressor(
+ binary_operators=["+", "*", "-", "/"],
+ logger_spec=logger_spec,
+)
+model.fit(X, y)
+```
+
+Start TensorBoard with the directory that contains the run logs:
+
+```bash
+tensorboard --logdir logs/
+```
+
+## Recording the genealogy of a search
+
+
+
+The complete recorder and post-processing workflow is in `examples/search_trace.py`. `load_trace` and `ancestry` are example-local helpers, not PySR API methods.
+
+Tracing records how a symbolic-regression population changes. Setting `use_tracing=True` activates the optional recorder. PySR loads the required JSON dependency lazily on first use. It then writes JSONL records to `tracing_file`, one for each population at each iteration, with the live members and their mutation, crossover, tuning, and death events. These recorded fields supply the members and parent relationships that the helpers need to reconstruct a graph and one selected path. The helpers work from recorded events and do not expose every historical branch.
+
+With tracing disabled, PySR does not request any trace output or load the tracing dependency.
+
+The target combines one quadratic term, one cosine term, and an offset. The added Gaussian noise has standard deviation equal to $5\%$ of the target's spread:
+
+```python
+import numpy as np
+
+rng = np.random.default_rng(0)
+X = rng.uniform(-3, 3, size=(200, 3))
+y = 2.5382 * np.cos(X[:, 2]) + X[:, 0] ** 2 - 1.5
+y = y + rng.normal(0, 0.05 * y.std(), size=200)
+```
+
+```python
+from pysr import PySRRegressor
+
+model = PySRRegressor(
+ binary_operators=["+", "-", "*", "/"],
+ unary_operators=["cos", "exp"],
+ population_size=27,
+ populations=8,
+ ncycles_per_iteration=90,
+ maxsize=20,
+ niterations=25,
+ random_state=0,
+ use_tracing=True,
+ tracing_file="search_trace.jsonl",
+)
+model.fit(X, y)
+```
+
+Keep tracing runs short: the files grow with the search, and this `load_trace` helper reads the entire file into memory.
+
+Data generation code
+ +Each scalar input is embedded as a diagonal `Force`, with the same value in each of the +three fields. The output remains a full three-component force. This lets the template recover +scalar inputs from `.x` while keeping one custom value type for both `X` and `y`: + +`TypeSpec` converts each tuple in these object arrays into the generated `Force` using the +declared field order (`x`, `y`, `z`), so Python does not need to call `Force(...)` directly. + +```python +import numpy as np + +from pysr import PySRRegressor, TemplateExpressionSpec, TypeSpec + +OMEGA = 2 * np.pi + + +def experiments(n, seed): + rng = np.random.default_rng(seed) + T = 298.15 + 0.5 * rng.random(n) + t = 10.0 * rng.random(n) + v = 2.0 * rng.random((n, 3)) - 1.0 + B = np.stack([np.sin(OMEGA * t), np.cos(OMEGA * t), np.exp(-t / 10.0)], axis=1) + F = -1e-5 * np.sqrt(T)[:, None] * v + np.cross(v, B) + return np.column_stack([t, v[:, 0], v[:, 1], v[:, 2], T]), F + + +def as_forces(rows): + """An (n, 3) float array as a column of n Force values.""" + out = np.empty(len(rows), dtype=object) + out[:] = [tuple(map(float, row)) for row in rows] + return out + + +def on_diagonal(columns): + out = np.empty(columns.shape, dtype=object) + for j in range(columns.shape[1]): + out[:, j] = as_forces(np.repeat(columns[:, j, None], 3, axis=1)) + return out + + +INPUTS, FORCES = experiments(1000, seed=0) +X = on_diagonal(INPUTS) +y = as_forces(FORCES) +variable_names = ["t", "v_x", "v_y", "v_z", "T"] + +HELD_OUT_INPUTS, HELD_OUT_FORCES = experiments(500, seed=12345) +HELD_OUT_X = on_diagonal(HELD_OUT_INPUTS) +``` + +
+
+
+`ancestry` selects the lowest-loss recorded member and follows the first edge for each child. It returns one path, which ends when no edge remains or a reference repeats.
+
+```python
+def ancestry(nodes, edges):
+ first_edge = {}
+ for edge in edges:
+ first_edge.setdefault(edge["child"], edge)
+ winner = min(
+ (ref for ref, n in nodes.items() if n.get("loss") is not None),
+ key=lambda ref: nodes[ref]["loss"],
+ )
+ chain, current, seen = [], winner, set()
+ while current in first_edge and current not in seen:
+ seen.add(current)
+ chain.append(first_edge[current])
+ current = first_edge[current]["parent"]
+ chain.reverse()
+ return winner, chain
+```
+
+Trace records label features `x1`, `x2`, and so on by column order, independently of any `variable_names` passed to the search. Thus `x1` refers to `X[:, 0]`, and `x3` refers to `X[:, 2]` in this example.
+
+Use `examples/search_trace.py` for the complete tracing workflow.
+
+## Closing an agent loop with `guesses=`
+
+The `guesses=` keyword accepts a list of expression strings to seed a fit. An external workflow can inspect `model.equations_`, propose an expression, and pass it into a later search. PySR can then optimize its numeric constants.
+
+This example replays a recorded coding-agent proposal from `PROPOSALS`. The custom `agent(front, tried)` helper returns stored expressions; it does not call a live agent or model service. The complete implementation is in `examples/agent_loop_guesses.py`.
+
+The target is relativistic kinetic energy in natural units, with $c=1$ and speeds up to $0.88c$. Gaussian noise has standard deviation $0.15\%$ of the median exact energy. `NOISE_FLOOR` stores its variance for the loop's training-loss stopping rule.
+
+Trace-reading code
+ +```python +import json +import pathlib + +def load_trace(trace_path): + nodes, edges = {}, [] + for line in pathlib.Path(trace_path).read_text().splitlines(): + record = json.loads(line) + if record["record_type"] != "iteration": + continue + for member in record["members"]: + nodes.setdefault(member["ref"], {}).update(member) + for ref, entry in record["mutations"].items(): + parent = int(ref) + nodes.setdefault(parent, {}).update( + {k: entry[k] for k in ("tree", "loss", "cost", "parent") if k in entry} + ) + for event in entry["events"]: + kind = event["type"] + detail = event.get("mutation") or event.get("details") or {} + if kind == "crossover" and event["parent1"] != parent: + continue + children = ( + (event["child1"], event["child2"]) + if kind == "crossover" + else (event["child"],) if "child" in event else () + ) + for child in children: + edges.append( + dict( + parent=parent, + child=child, + kind=kind, + mutation=detail.get("type"), + result=detail.get("result", "accept"), + mate=event.get("parent2"), + ) + ) + for edge in edges: + if ( + edge["mate"] is not None + and nodes.get(edge["child"], {}).get("parent") == edge["mate"] + ): + edge["parent"], edge["mate"] = edge["mate"], edge["parent"] + return nodes, edges +``` + +
+
+
+```python
+from pysr import PySRRegressor
+
+model_kwargs = dict(
+ operators={1: ["sqrt"], 2: ["+", "-", "*", "/"]},
+ niterations=10,
+ verbosity=0,
+)
+```
+
+The proposal uses the known relativistic structure:
+
+```python
+PROPOSALS = ["(mass / sqrt(1 - (velocity * velocity))) - mass"]
+
+def agent(front, tried):
+ print("front handed to the agent:")
+ print(
+ front[["complexity", "loss", "equation"]].to_string(index=False),
+ "\n",
+ flush=True,
+ )
+ return PROPOSALS[len(tried)] if len(tried) < len(PROPOSALS) else None
+```
+
+The loop begins with `guesses=None` and passes each proposal as a one-element list to the next fit. It stops when the minimum training loss reaches `2 * NOISE_FLOOR` or the helper has no proposal left. This threshold is an example-specific stopping rule, not a lower bound on achievable loss.
+
+```python
+from pysr import PySRRegressor
+
+guesses, tried = None, []
+while True:
+ model = PySRRegressor(**model_kwargs, guesses=guesses, random_state=0)
+ model.fit(X, y, variable_names=VARIABLE_NAMES)
+ front = model.equations_
+ if front["loss"].min() <= 2 * NOISE_FLOOR:
+ break
+ proposal = agent(front, tried)
+ if proposal is None:
+ break
+ tried.append(proposal)
+ guesses = [proposal]
+```
+
+Use `examples/agent_loop_guesses.py` for the complete handoff workflow.
diff --git a/docs/src/examples/objectives.md b/docs/src/examples/objectives.md
new file mode 100644
index 000000000..1766b33c8
--- /dev/null
+++ b/docs/src/examples/objectives.md
@@ -0,0 +1,332 @@
+# Objectives and losses
+
+## Custom objectives
+
+`loss_function` gives the objective the candidate's expression tree. Use it when the score needs to manipulate the tree or read dataset-level information. Use `loss_function_expression` when the score needs the complete expression object. Both modes accept Julia source with one of the signatures below.
+
+```julia
+objective(tree_or_expression, dataset, options)
+objective(tree_or_expression, dataset, options, idx=nothing)
+```
+
+With three arguments, the objective sees the dataset given to the current evaluation. Automatic batching makes that dataset the active batch. A method with `idx=nothing` can receive the full dataset and the selected row indices, allowing the objective to slice the rows itself.
+
+The following objective reproduces ordinary unweighted mean squared error, while retaining access to the optional batch indices:
+
+```python
+from pysr import PySRRegressor
+
+objective = """
+function mse_objective(tree, dataset::Dataset{T,L}, options, idx=nothing)::L where {T,L}
+ X = idx === nothing ? dataset.X : dataset.X[:, idx]
+ y = idx === nothing ? dataset.y : dataset.y[idx]
+ prediction, complete = eval_tree_array(tree, X, options)
+ complete || return L(Inf)
+ return sum(abs2, prediction .- y) / length(y)
+end
+"""
+
+model = PySRRegressor(
+ loss_function=objective,
+ binary_operators=["+", "-", "*", "/"],
+)
+```
+
+`eval_tree_array` returns a completion flag with the predictions. Check that flag before using the array and assign an infinite, or suitably large, loss to an incomplete evaluation. Return the dataset's loss type `L`; it must remain real even when the data type `T` is complex.
+
+PySR imports `SymbolicRegression` into Julia before evaluating these source strings, making `Dataset` and `eval_tree_array` available without another import.
+
+A full objective can also attach a different meaning to the stored tree. Here the two children of a binary root supply the numerator and denominator, forming the rational function $P(X)/Q(X)$. The fully qualified `DE.get_child` calls below access the backend's tree interface.
+
+```python
+from pysr import PySRRegressor
+
+objective = """
+function rational_objective(tree, dataset::Dataset{T,L}, options, idx=nothing)::L where {T,L}
+ tree.degree == 2 || return L(Inf)
+ X = idx === nothing ? dataset.X : dataset.X[:, idx]
+ y = idx === nothing ? dataset.y : dataset.y[idx]
+
+ numerator = SymbolicRegression.InterfaceDynamicExpressionsModule.DE.get_child(tree, 1)
+ denominator = SymbolicRegression.InterfaceDynamicExpressionsModule.DE.get_child(tree, 2)
+ p, p_complete = eval_tree_array(numerator, X, options)
+ q, q_complete = eval_tree_array(denominator, X, options)
+ p_complete && q_complete || return L(Inf)
+
+ prediction = p ./ q
+ return sum(abs2, prediction .- y) / length(y)
+end
+"""
+
+model = PySRRegressor(
+ loss_function=objective,
+ binary_operators=["+", "-"],
+)
+```
+
+This objective ignores the root operator and interprets only its two children. The equation table displays the stored tree. `predict` and symbolic export likewise retain the raw tree and cannot automatically reconstruct the quotient. Keep the objective source and apply the child transformation yourself whenever you evaluate the selected equation.
+
+## Writing the objective in Python
+
+The objective itself can live in Python while Julia supplies a thin shim. The search releases Python's global interpreter lock (GIL) while it runs, so the shim reacquires the lock, calls the Python function, and converts its scalar result to the dataset's loss type. The Python function receives the tree, dataset, and options as Julia objects, and it can call back into Julia, including `eval_tree_array`:
+
+```python
+import numpy as np
+
+from pysr import PySRRegressor, jl
+
+
+def python_objective(tree, dataset, options):
+ prediction, completed = jl.SymbolicRegression.eval_tree_array(
+ tree, dataset.X, options
+ )
+ if not completed:
+ return float("inf")
+
+ prediction = np.asarray(prediction)
+ target = np.asarray(dataset.y)
+ return float(np.mean((prediction - target) ** 2))
+
+
+jl.python_objective = python_objective
+
+jl.seval("""
+using PythonCall
+
+function python_objective_shim(
+ tree, dataset::Dataset{T,L}, options
+)::L where {T,L}
+ PythonCall.GIL.@lock begin
+ return pyconvert(L, python_objective(tree, dataset, options))
+ end
+end
+""")
+
+model = PySRRegressor(
+ loss_function="python_objective_shim",
+ precision=64,
+)
+```
+
+`precision=64` makes the data and loss use `Float64`, matching the Python float returned by this objective.
+
+Use this bridge when the objective needs Python-specific logic. Every evaluation crosses the Julia/Python boundary, and concurrent callbacks serialize on the GIL. Keep throughput-sensitive objectives in Julia.
+
+## Swinging up a cart-pole with a rollout objective
+
+
+
+A supervised example gives a target for each row. A control policy has no label of that kind for its force: a force is good only if the subsequent trajectory is good. `loss_function_expression` provides the candidate expression to a rollout objective that can evaluate it throughout the closed loop. This example runs 500 steps from each of 16 training starts and returns the negative mean reward per step. Most of the starts leave the pole hanging down, so the expression must swing it up, catch it, and balance it with one closed-form policy.
+
+### The plant and what the policy sees
+
+The implemented plant uses a $1.0\,\mathrm{kg}$ cart and a $0.1\,\mathrm{kg}$ pole with half-length $0.5\,\mathrm{m}$, gravity $9.8$, and semi-implicit Euler integration with $dt=0.02\,\mathrm{s}$. Its 500-step horizon is therefore $10\,\mathrm{s}$. The policy output is clipped to `[-1, 1]` and multiplied by $10\,\mathrm{N}$. The five inputs are scaled cart position, scaled cart speed, $\sin\theta$, $\cos\theta$, and scaled pole rate. The sine/cosine pair avoids the wrap discontinuity at $\pm\pi$, while the cosine input `c` supplies a direct uprightness signal.
+
+Data generation code
+ +```python +import numpy as np + +rng = np.random.default_rng(42) +mass = rng.uniform(0.5, 3.0, 160) +velocity = rng.uniform(0.08, 0.88, 160) +kinetic_energy_exact = mass * (1 / np.sqrt(1 - velocity**2) - 1) +NOISE_SIGMA = 0.0015 * float(np.median(kinetic_energy_exact)) +kinetic_energy = kinetic_energy_exact + rng.normal(0.0, NOISE_SIGMA, mass.size) +NOISE_FLOOR = NOISE_SIGMA**2 + +X = np.stack([mass, velocity], axis=1) +y = kinetic_energy +VARIABLE_NAMES = ["mass", "velocity"] +``` + +
+
+
+Use the initial-state arrays and `observe` helper in `examples/cartpole_objective.py` to construct `X`. Its columns contain the five policy observations, and `y` is a zero placeholder that the objective ignores. The 16 training starts cover hanging, horizontal, and near-upright states; the 64 held-out starts use wider ranges of position, speed, angle, and pole rate.
+
+### The reward
+
+At each step, the objective accumulates the reward
+
+$$ r = 2\cos\theta - 0.05\,x^2 - 0.01\left(\frac{F}{10}\right)^2 - 0.005\left(\frac{F - F_{\text{prev}}}{10}\right)^2 $$
+
+The upright term contributes up to 2 per step. The remaining terms penalize, in order, cart displacement, normalized actuator force, and normalized force changes. Here $F$ is the clipped force and $F_{\mathrm{prev}}$ is the preceding force. The objective returns the negative mean of $r$ over 500 steps and 16 starts, so an upright stationary policy has loss near $-2$.
+
+The complete objective source is defined in `examples/cartpole_objective.py` as the Python string `CARTPOLE_REWARD`. The block below shows its Julia function body and abbreviates the state unpacking, observation-buffer setup, integration step, and divergence check supplied by that script. Treat it as a behavior excerpt, not a standalone definition.
+
+```julia
+function cartpole_reward(ex, dataset::Dataset{T,L}, options)::L where {T,L}
+ n = size(dataset.X, 2)
+ # ... unpack x, v, theta, omega from the observation columns ...
+ for _ in 1:500
+ # ... refill obs from the current state ...
+ raw, ok = eval_tree_array(ex, obs, options)
+ (ok && length(raw) == n && all(isfinite, raw)) || return L(Inf)
+
+ force = 10.0 .* clamp.(raw, -1.0, 1.0)
+ earned += sum(
+ 2.0 .* cos_t .- 0.05 .* x .^ 2 .-
+ 0.01 .* (force ./ 10.0) .^ 2 .-
+ 0.005 .* ((force .- previous) ./ 10.0) .^ 2
+ )
+ previous = force
+ # ... one semi-implicit Euler step, then the divergence check ...
+ end
+ return L(-earned / (500 * n))
+end
+```
+
+Invalid evaluations and divergent trajectories receive a penalty.
+
+### The search
+
+```python
+import sympy
+
+from pysr import PySRRegressor
+
+model = PySRRegressor(
+ operators={
+ 1: ["square", "abs", "tanh"],
+ 2: ["+", "-", "*", "/", "max", "min"],
+ 3: ["ifelse(t, a, b) = t > 0 ? a : b"],
+ },
+ extra_sympy_mappings={
+ "ifelse": lambda t, a, b: sympy.Piecewise((a, t > 0), (b, True))
+ },
+ loss_function_expression=CARTPOLE_REWARD,
+ loss_scale="linear", # rewards make the loss negative, which log scaling forbids
+ maxsize=28,
+ maxdepth=10,
+ parsimony=0.001,
+ niterations=120,
+ populations=32,
+ population_size=64,
+ random_state=0,
+)
+model.fit(X, y, variable_names=["x_n", "v_n", "s", "c", "omega_n"])
+```
+
+`loss_scale="linear"` is required here because logarithmic scaling accepts nonnegative losses, including zero, while this reward can make the loss negative. The operator set supplies arithmetic, `max`, and `min`; the user-defined ternary `ifelse` can switch between swing-up and balancing behavior. `extra_sympy_mappings` translates that custom operator to a SymPy `Piecewise` expression for export. The search uses multithreading, so repeated fits can differ even with the same nominal seed.
+
+### Evaluating a policy
+
+One controller returned by the search is:
+
+```
+(tanh(((0.46341985 * v_n) + (omega_n + (s + omega_n))) / 0.00043365502) - v_n) / 0.97654253
+```
+
+A low training loss alone does not establish successful swing-up from new initial states. In `examples/cartpole_objective.py`, `mean_rewards(model, HELD_OUT_STARTS, index=...)` evaluates a selected equation on the wider held-out set, and `check(model)` requires positive mean reward on every start. Use these rollout checks when choosing a policy.
+
+Allow roughly two hours for the search configuration above. The complete runnable example is `examples/cartpole_objective.py`.
+
+## Inventing a pseudorandom generator with no target
+
+
+
+Ordinary regression supplies `y` and compares every candidate output with a target. This example instead searches for a 32-bit state update $x \mapsto f(x)$ whose iterates have specified properties. Its score comes from how the map behaves, so the objective defines what counts as a good generator without consulting a target array.
+
+The data rows provide 32 random nonzero seeds. `y` is a zero column retained for the estimator interface; the objective never reads it.
+
+```python
+import numpy as np
+
+from pysr import PySRRegressor, TypeSpec
+
+WIDTH = 32
+SEEDS = [int(v) for v in np.random.default_rng(0).integers(1, 1 << WIDTH, size=32)]
+X = np.empty((len(SEEDS), 1), dtype=object)
+for i, seed in enumerate(SEEDS):
+ X[i, 0] = seed
+y = np.zeros(len(SEEDS), dtype=object)
+```
+
+Each state is an unsigned 32-bit word, wrapped in a Julia type. The `shl` and `shr` operators below pass their `UInt32` shift count directly; this example does not mask counts into `0:31`, so arbitrary sampled counts follow Julia's shift semantics and bits shifted out of the 32-bit word are discarded. The operations do not rotate the state. The `sample` hook favors values from 0 through 31 because those are common shift distances, while still sampling arbitrary `UInt32` values. Half of mutations flip a single bit, allowing a local change to a tap constant; the other half adds either 1 or `typemax(UInt32)`. The `string` hook displays values below 32 in decimal and larger words as eight hexadecimal digits.
+
+```python
+from pysr import TypeSpec
+
+SPEC = TypeSpec(
+ "Word",
+ fields={"bits": "UInt32"},
+ sample="rng -> Word(rand(rng, Bool) ? rand(rng, UInt32(0):UInt32(31)) : rand(rng, UInt32))",
+ mutate="(rng, value, temperature) -> Word(rand(rng, Bool) ? xor(value.bits, UInt32(1) << rand(rng, 0:31)) : value.bits + rand(rng, (UInt32(1), typemax(UInt32))))",
+ string='value -> value.bits < UInt32(32) ? string(value.bits) : "0x" * string(value.bits, base = 16, pad = 8)',
+ loss_type="Float64",
+)
+```
+
+The operators form a restricted bitwise instruction set: complement, XOR, AND, OR, left shift, and right shift. The search excludes arithmetic and floating-point operations.
+
+### What the objective scores
+
+`examples/prng_period.py` defines the objective as a Julia source string named `PRNG_LOSS`. It includes the period-certificate implementation used below.
+
+For every candidate, the objective walks 256 steps from each seed and adds six terms that penalize departures from their desired properties:
+
+- `period` is $1 - \log_2(\mathrm{proven})/32$, with `proven` equal to the certified orbit length when certification succeeds.
+- `balance` is the mean squared bias of each bit column away from one half.
+- `autocorr` is the mean squared normalized self-correlation of each bit column at lags 1, 2, 3, and 5.
+- `diffusion` penalizes an average of fewer than four changed output bits when each input bit is flipped on a fixed probe set.
+- `edge` penalizes output bits that respond to fewer than two distinct input-bit flips across that probe set.
+- `shear` measures correlations between bit `b` and shifted bit `b + d` at one- and two-step lags, for nonzero shifts `d` from -3 to 3. It keeps a plain shift from receiving a good score.
+
+### Certifying a period of four billion without walking it
+
+The requested nonzero orbit contains $2^{32}-1 = 4{,}294{,}967{,}295$ states, so a direct walk is impractical. The certificate evaluates the candidate at three groups of inputs: zero, all 32 basis states, and 32 fixed probes. It requires zero to map to zero and the image of each probe to agree with the matrix assembled from the basis images. These finite checks only screen for GF(2) linearity on the supplied evidence; they do not prove global linearity for an arbitrary expression.
+
+A successful certificate gives the candidate's period provided the expression is globally GF(2)-linear. If certification fails, the objective uses the orbit length observed during the finite walk. The period-checking implementation is in `examples/prng_period.py`.
+
+### Settings
+
+The settings below pass the source-defined `PRNG_LOSS` objective to PySR:
+
+```python
+from pysr import PySRRegressor
+
+model = PySRRegressor(
+ type_spec=SPEC,
+ operators={
+ 1: ["bnot(a::Word) = Word(~a.bits)"],
+ 2: [
+ "bxor(a::Word, b::Word) = Word(xor(a.bits, b.bits))",
+ "band(a::Word, b::Word) = Word(a.bits & b.bits)",
+ "bor(a::Word, b::Word) = Word(a.bits | b.bits)",
+ "shl(a::Word, b::Word) = Word(a.bits << b.bits)",
+ "shr(a::Word, b::Word) = Word(a.bits >> b.bits)",
+ ],
+ },
+ loss_function=PRNG_LOSS,
+ batching=False,
+ maxsize=20,
+ populations=8,
+ population_size=30,
+ ncycles_per_iteration=30,
+ niterations=800,
+ random_state=0,
+ verbosity=0,
+)
+model.fit(X, y, variable_names=["x"])
+```
+
+Keep `batching=False` because the objective walks every seed; a minibatch would score a different problem. Each candidate requires a 256-step walk over all 32 seeds. Allow about an hour for this search.
+
+### Results
+
+The script's `check(model)` looks for a front member that passes a separate finite linearity screen and whose inferred GF(2) matrix has order $2^{32}-1$. As with the objective's certificate, this establishes the candidate's period only if the expression is globally GF(2)-linear.
+
+One recovered expression is:
+
+```
+bxor(shl(x, 1), bxor(shr(bxor(shr(x, 1), x), 1), band(x, 0xffffffe3)))
+```
+
+Fixed shifts, XOR, and AND with a constant mask make this expression globally linear by construction. For such a map, matrix order $2^{32}-1$ implies that every nonzero word lies on one orbit of that length. The remaining score terms assess the generator's finite-sample balance, correlations, and diffusion.
+
+The complete runnable example is `examples/prng_period.py`.
diff --git a/docs/src/examples/physics.md b/docs/src/examples/physics.md
new file mode 100644
index 000000000..156c66f27
--- /dev/null
+++ b/docs/src/examples/physics.md
@@ -0,0 +1,183 @@
+# Physics and units
+
+## Dimensional constraints
+
+Dimensional constraints let PySR use physical units during symbolic search. When the units of the features and target are known, pass a unit expression for each feature and one for the target so candidates with mismatched dimensions receive a penalty.
+
+The fit interface follows the unit syntax of `DynamicQuantities.jl`. The following example uses Astropy to sample a mass $M$, a test mass $m$, and a separation $r$, then evaluates Newton's inverse-square law $F = GMm/r^2$.
+
+```python
+import numpy as np
+from astropy import units as u, constants as const
+
+M = (np.random.rand(100) + 0.1) * const.M_sun
+m = 100 * (np.random.rand(100) + 0.1) * u.kg
+r = (np.random.rand(100) + 0.1) * const.R_earth
+G = const.G
+
+F = G * M * m / r**2
+```
+
+The force values span a wide range. This loss compares their magnitudes in logarithmic space and penalizes a wrong sign:
+
+```python
+elementwise_loss = """function loss_fnc(prediction, target)
+ scatter_loss = abs(log((abs(prediction)+1f-20) / (abs(target)+1f-20)))
+ sign_loss = 10 * (sign(prediction) - sign(target))^2
+ return scatter_loss + sign_loss
+end
+"""
+```
+
+`dimensional_constraint_penalty` adds a cost when a candidate fails dimensional analysis:
+
+```python
+from pysr import PySRRegressor
+
+model = PySRRegressor(
+ binary_operators=["+", "-", "*", "/"],
+ unary_operators=["square"],
+ elementwise_loss=elementwise_loss,
+ complexity_of_constants=2,
+ maxsize=25,
+ niterations=100,
+ populations=50,
+ # Amount to penalize dimensional violations:
+ dimensional_constraint_penalty=10**5,
+)
+```
+
+Pass numeric values to `fit`, with matching units in column order through `X_units` and the target unit through `y_units`. Unit strings follow the Julia syntax documented by [DynamicQuantities.jl](https://ai.damtp.cam.ac.uk/dynamicquantities/dev/#Usage).
+
+```python
+import pandas as pd
+
+# Get numerical arrays to fit:
+X = pd.DataFrame(dict(
+ M=M.to("M_sun").value,
+ m=m.to("kg").value,
+ r=r.to("R_earth").value,
+))
+y = F.value
+
+model.fit(
+ X,
+ y,
+ X_units=["Constants.M_sun", "kg", "Constants.R_earth"],
+ y_units="kg * m / s^2"
+)
+```
+
+The loss is nonnegative, and the dimensional penalty is additive. A candidate with total loss below that penalty therefore did not receive the dimensional-violation charge and is dimensionally consistent for the supplied units. The marker `"[⋅]"` denotes a constant whose units remain free. The unit-aware form below illustrates this notation: `M[kg]` supplies mass, so the constant supplies the acceleration units needed for the target:
+
+```julia
+"y[m s⁻² kg] = (M[kg] * 2.6353e-22[⋅])"
+```
+
+Set `dimensionless_constants_only=True` to restrict fitted constants to dimensionless quantities.
+
+## Using differential operators
+
+A [`TemplateExpressionSpec`](/examples/expression-specifications) can include differential operators in its `combine` string. `D` takes the expression to differentiate and an argument index. Here, `D(f, 1)` differentiates the learned function `f` with respect to its first argument, and `df(x)` evaluates that derivative. Matching those values to data turns antiderivative discovery into regression.
+
+The regression target is the integrand $1/(x^2\sqrt{x^2 - 1})$ on $x > 1$. The template searches for a function $f$ whose derivative matches samples of the integrand, leaving an additive constant undetermined.
+
+```python
+import numpy as np
+
+from pysr import PySRRegressor, TemplateExpressionSpec
+
+x = np.random.uniform(1, 10, (1000,)) # Integrand sampling points
+y = 1 / (x**2 * np.sqrt(x**2 - 1)) # Evaluation of the integrand
+
+expression_spec = TemplateExpressionSpec(
+ expressions=["f"],
+ variable_names=["x"],
+ combine="df = D(f, 1); df(x)",
+)
+
+model = PySRRegressor(
+ binary_operators=["+", "-", "*", "/"],
+ unary_operators=["sqrt"],
+ expression_spec=expression_spec,
+ maxsize=20,
+)
+model.fit(x[:, np.newaxis], y)
+```
+
+The target antiderivative is $f(x) = \frac{\sqrt{x^2 - 1}}{x}$, up to an additive constant.
+
+## Discovering a PDE
+
+For a field $u(x,t)$ sampled on a space-time grid, each point becomes one regression row: its field value and spatial derivatives are features, and its time derivative is the target. PySR then searches for the right-hand side in $u_t = f(u, u_x, u_{xx}, \ldots)$.
+
+The following example simulates the viscous Burgers equation $u_t = -u\,u_x + 0.1\,u_{xx}$ on a periodic domain. Measurements of the field can replace the simulation when they are arranged into the same feature and target arrays.
+
+Constants and the plant step
+ +```python +import numpy as np + +M_C, M_P, LENGTH, GRAVITY = 1.0, 0.1, 0.5, 9.8 +FORCE_CAP, DT = 10.0, 0.02 +RAIL, V_SCALE, OMEGA_SCALE = 2.4, 4.0, 8.0 +HORIZON = 500 # 10 s of control at dt = 0.02 + + +def step(state, force): + x, v, theta, omega = state.T + total = M_C + M_P + sin, cos = np.sin(theta), np.cos(theta) + q = (force + M_P * LENGTH * omega**2 * sin) / total + theta_dd = (GRAVITY * sin - cos * q) / (LENGTH * (4 / 3 - M_P * cos**2 / total)) + x_dd = q - M_P * LENGTH * theta_dd * cos / total + v, omega = v + DT * x_dd, omega + DT * theta_dd + return np.stack([x + DT * v, v, theta + DT * omega, omega], axis=-1) +``` + +
+
+
+Estimate spatial derivatives with the smoothing differentiator `savgol_filter` and time derivatives with `np.gradient`. Supply the coordinate spacings and keep each space-time point aligned across the feature columns and target:
+
+```python
+import numpy as np
+
+from scipy.signal import savgol_filter
+
+ux = savgol_filter(U, 21, 3, deriv=1, delta=dx, axis=-1)
+uxx = savgol_filter(U, 21, 3, deriv=2, delta=dx, axis=-1)
+ut = np.gradient(U, t, axis=0)
+
+X = np.stack([U, ux, uxx], axis=-1).reshape(-1, 3)
+y = ut.reshape(-1)
+```
+
+`complexity_of_variables=[1, 2, 3]` makes higher spatial derivatives more expensive, biasing the search toward lower-order terms when candidate fits compete. A successful low-complexity search should place $0.1\,u_{xx} - u\,u_x$ on the Pareto front, matching the right-hand side that generated the snapshots.
+
+```python
+from pysr import PySRRegressor
+
+model = PySRRegressor(
+ binary_operators=["+", "-", "*"],
+ complexity_of_variables=[1, 2, 3],
+ maxsize=20,
+ niterations=100,
+)
+model.fit(X, y, variable_names=["u", "u_x", "u_xx"])
+print(model)
+```
+
+The regression uses estimated derivatives, so simulate the discovered PDE forward from a held-out initial condition before relying on its coefficients.
+
+If the field $u$ has units $\mathrm{m\,s^{-1}}$, then $u_x$ has units $\mathrm{s^{-1}}$, $u_{xx}$ has units $\mathrm{m^{-1}\,s^{-1}}$, and $u_t$ has units $\mathrm{m\,s^{-2}}$. For this feature order, pass the corresponding strings through `X_units=["m/s", "s^-1", "m^-1*s^-1"]` and `y_units="m*s^-2"`.
diff --git a/docs/src/examples/search-behaviour.md b/docs/src/examples/search-behaviour.md
new file mode 100644
index 000000000..e46f0cb7f
--- /dev/null
+++ b/docs/src/examples/search-behaviour.md
@@ -0,0 +1,270 @@
+# Search behaviour
+
+## Automatic batching on a large dataset
+
+`batching` reduces the rows used to score population members during evolution. Hall-of-fame comparisons continue to use every row. With `batching=False`, every evolutionary comparison uses all rows. When `batch_size` is `None`, automatic batching can use every row on small datasets. PySR chooses all rows through 1000, 128 below 5000, 256 below 50,000, and 512 from 50,000 onward. At 20,000 rows, evolution therefore sees 256 rows per comparison, while hall-of-fame scoring sees all 20,000. `batching="auto"` is the default; this example names it to make the choice explicit.
+
+The target depends on two of five input columns. The remaining three inputs act as distractors.
+
+Data generation code
+ +```python +import numpy as np + +L, nx, nu = 2 * np.pi, 128, 0.1 +x = np.linspace(0.0, L, nx, endpoint=False) +dx = L / nx +k = 2 * np.pi * np.fft.rfftfreq(nx, d=dx) + +def d_dx(u, order=1): + return np.real(np.fft.irfft((1j * k) ** order * np.fft.rfft(u), n=nx)) + +def rhs(u): + return -u * d_dx(u) + nu * d_dx(u, order=2) + +u = -np.sin(x) +snapshots = [u.copy()] +dt = 1e-3 +for i in range(1, 4001): + k1 = rhs(u) + k2 = rhs(u + 0.5 * dt * k1) + k3 = rhs(u + 0.5 * dt * k2) + k4 = rhs(u + dt * k3) + u = u + (dt / 6.0) * (k1 + 2 * k2 + 2 * k3 + k4) + if i % 100 == 0: + snapshots.append(u.copy()) + +U = np.array(snapshots) # shape (41, 128) +t = np.arange(U.shape[0]) * 0.1 +``` + +
+
+
+```python
+from pysr import PySRRegressor
+
+model = PySRRegressor(
+ binary_operators=["+", "-", "*", "/"],
+ unary_operators=["cos", "exp"],
+ batching="auto",
+ niterations=40,
+ verbosity=0,
+ random_state=0,
+)
+model.fit(X, y, variable_names=["x0", "x1", "x2", "x3", "x4"])
+print(model.equations_[["complexity", "loss", "equation"]].to_string(index=False))
+```
+
+One recovered expression is:
+
+```
+(cos(x3) * 2.5382) + ((x0 * x0) + -0.5)
+```
+
+Run the complete example with `examples/automatic_batching.py`.
+
+## Operators of any arity
+
+
+
+The generic `operators` argument defines an expression grammar by arity. Each integer key gives the number of arguments accepted by the operators listed under it, so one search can mix unary and binary operations with ternary or wider functions. `operators` cannot be combined with `binary_operators` or `unary_operators`.
+
+The target is a four-way minimum,
+
+$$
+y = \min(x + 2,\; 2 - x,\; \sin(3x) + 1.2,\; 0.8).
+$$
+
+A binary `min` would write this reduction as three nested calls. The custom `min4` writes the same mathematical operation in one call, giving the search a shorter structural route to the target.
+
+Data generation code
+ +```python +import numpy as np + +rng = np.random.default_rng(0) +X = rng.uniform(-3, 3, (20000, 5)) +y = 2.5382 * np.cos(X[:, 3]) + X[:, 0] ** 2 - 0.5 +``` + +
+
+
+Custom operator definitions are Julia source strings under their arity keys. A SymPy-capable model also needs a matching entry in `extra_sympy_mappings`; PySR rejects a named custom function during setup when its symbolic image is missing. The mappings here provide the `Piecewise` representation for `ifelse` and the `Min` representation for `min4`.
+
+```python
+import sympy
+
+from pysr import PySRRegressor
+
+model = PySRRegressor(
+ operators={
+ 1: ["sin", "cos"],
+ 2: ["+", "-", "*"],
+ 3: ["ifelse(t, a, b) = t > 0 ? a : b"],
+ 4: ["min4(a, b, c, d) = min(a, b, c, d)"],
+ },
+ # A custom operator is rejected at fit time unless it has a SymPy image.
+ extra_sympy_mappings={
+ "ifelse": lambda t, a, b: sympy.Piecewise((a, t > 0), (b, True)),
+ "min4": lambda a, b, c, d: sympy.Min(a, b, c, d),
+ },
+ niterations=500,
+ random_state=0,
+)
+model.fit(X, y, variable_names=["x"])
+print(model.equations_)
+```
+
+One recovered expression is:
+
+```
+min4(0.8, 2.0 - x, 2.0 + x, 1.1999999 + sin(x * 3.0))
+```
+
+The complete example, `examples/any_arity.py`, checks predictions on an unseen grid with `model.predict(X_UNSEEN, index=i)` for each candidate equation.
+
+## Mutations and plugins
+
+PySR accepts mutation and plugin configuration as `PySRRegressor` arguments. `mutations` maps mutation entries to weights, and `plugins` lists plugin entries. Entries merge with shipped defaults by concrete type: a matching type replaces the default entry, while a new type is appended alongside it. This lets one behavior change without reproducing the complete default set.
+
+The following configuration activates Backsolve and names the adaptive mutation-weight plugin.
+
+```python
+from pysr import (
+ AdaptiveMutationWeightsPlugin,
+ BacksolveMutation,
+ PySRRegressor,
+)
+
+model = PySRRegressor(
+ mutations={BacksolveMutation(): 0.1},
+ plugins=[AdaptiveMutationWeightsPlugin()],
+)
+```
+
+`BacksolveMutation` is present at default weight 0.0, so it is available without being sampled. Assigning weight 0.1 makes it eligible for this search. Backsolve inverts the operators above a selected subtree and fits a replacement from expressions already in the population, helping an additive target assemble from separately discovered pieces.
+
+`AdaptiveMutationWeightsPlugin` is part of the default plugin set. It updates mutation weights from observed improvement rates and resets its runtime state for each `fit`. Naming it in `plugins` replaces the existing entry of that type, which is also where `smoothing`, `floor`, and `reward` can be changed. The resolved list keeps simulated annealing and adaptive parsimony alongside adaptive mutation weights. The two settings work together: the mutation weight makes Backsolve eligible, and the adaptive plugin can change how much of the mutation budget each kind receives.
+
+The runnable configuration uses the small target
+
+$$ y = 2.5\cos(3x) + 0.5x^2 - 1. $$
+
+```python
+import numpy as np
+
+x = np.linspace(-3.0, 3.0, 200)
+X = x.reshape(-1, 1)
+y = 2.5 * np.cos(3.0 * x) + 0.5 * x * x - 1.0
+```
+
+```python
+from pysr import AdaptiveMutationWeightsPlugin, BacksolveMutation, PySRRegressor
+
+model = PySRRegressor(
+ binary_operators=["+", "-", "*"],
+ unary_operators=["cos", "exp"],
+ niterations=20,
+ verbosity=0,
+ mutations={BacksolveMutation(): 0.1},
+ plugins=[AdaptiveMutationWeightsPlugin()],
+ random_state=0,
+)
+model.fit(X, y, variable_names=["x"])
+print(model.equations_)
+```
+
+After fitting, inspect `model.julia_options_.mutations` to see the resolved entries and weights. Run the example with `examples/mutations_and_plugins.py`.
+
+## Adaptive mutation weights
+
+
+
+`AdaptiveMutationWeightsPlugin` adjusts mutation weights from their observed success rates. It is enabled by default, and its learned state resets on each `fit`. Pass an instance through `plugins` to change its settings.
+
+- `smoothing` controls the exponential moving average of success rates.
+- `floor` limits the success-rate ratio used to update a multiplier to $[\text{floor}, 1/\text{floor}]$.
+- `reward` selects whether improvement is measured by `"cost"` or `"loss"`.
+
+This example fits a cosine and a quadratic term:
+
+Data generation code
+ +```python +import numpy as np + +DOMAIN = (-4.0, 4.0) + + +def envelope(x): + return np.minimum.reduce( + [x + 2.0, 2.0 - x, np.sin(3.0 * x) + 1.2, np.full_like(x, 0.8)] + ) + + +_x = np.sort(np.random.default_rng(0).uniform(*DOMAIN, 400)) +X = _x.reshape(-1, 1).astype(np.float32) +y = envelope(_x).astype(np.float32) + +_unseen = np.linspace(*DOMAIN, 4001) +X_UNSEEN = _unseen.reshape(-1, 1).astype(np.float32) +Y_UNSEEN = envelope(_unseen) +``` + +
+
+
+```python
+from pysr import AdaptiveMutationWeightsPlugin, PySRRegressor
+
+model = PySRRegressor(
+ binary_operators=["+", "-", "*", "/"],
+ unary_operators=["cos", "exp"],
+ maxsize=14,
+ niterations=300,
+ populations=1,
+ plugins=[AdaptiveMutationWeightsPlugin(smoothing=0.02, floor=0.05, reward="cost")],
+ random_state=0,
+ verbosity=0,
+)
+model.fit(X, y)
+print(model.equations_)
+```
+
+For an optional diagnostic that records mutation counts and learned multipliers, see `examples/adaptive_mutation_weights.py`.
+
+## The backsolve mutation
+
+
+
+Ordinary mutations change an expression and test the result. Backsolve chooses a non-root subtree and inverts the operators above it to determine the values that subtree would need to produce. It then fits a sparse replacement from expressions already in the population. Greedy forward selection over a basis library finds the replacement within the remaining complexity budget. The replaced subtree can remain in the library, and rows with invalid inverse values are masked. Both `+` and `*` are required for the weighted sum. An event returns a child only when its cost strictly improves on the parent. Each event requires an inversion and a linear solve, so enabling Backsolve can substantially increase search time.
+
+Backsolve applies to ordinary binary scalar expression trees with a real target vector. It rejects unsupported wrappers, shared nodes, root or single-node choices, missing targets, and missing weighted-sum operators. The mutation is experimental and starts at default weight 0.0.
+
+The target combines an oscillation with a quadratic:
+
+$$ y = 2.5\cos(3x) + 0.5x^2 - 1. $$
+
+```python
+import numpy as np
+
+x = np.linspace(-3.0, 3.0, 200)
+X = x.reshape(-1, 1)
+y = 2.5 * np.cos(3.0 * x) + 0.5 * x * x - 1.0
+```
+
+The example omits division and a power operator, so the quadratic must appear as `x * x`. `precision=64` supports the linear solve, and `maxsize=30` leaves room for the complete form.
+
+```python
+from pysr import PySRRegressor
+
+model = PySRRegressor(
+ binary_operators=["+", "-", "*"],
+ unary_operators=["cos", "exp"],
+ weight_backsolve=1.0,
+ maxsize=30,
+ precision=64,
+ niterations=10,
+ verbosity=0,
+ random_state=0,
+)
+model.fit(X, y, variable_names=["x"])
+print(model.equations_[["complexity", "loss", "equation"]].to_string(index=False))
+```
+
+One recovered expression is:
+
+```
+((cos((x + x) + x) * 2.5) + ((x * x) * 0.5)) + -1.0
+```
+
+Check the selected equation against the training target:
+
+```python
+import numpy as np
+
+prediction = np.asarray(model.predict(X), dtype=float)
+print(f"max abs error: {np.max(np.abs(prediction - y)):.3g}")
+```
+
+Use `examples/backsolve_mutation.py` for the complete example.
diff --git a/docs/src/examples/value-types.md b/docs/src/examples/value-types.md
new file mode 100644
index 000000000..53b679410
--- /dev/null
+++ b/docs/src/examples/value-types.md
@@ -0,0 +1,413 @@
+# Value types
+
+## Complex numbers
+
+To fit complex-valued expressions, pass feature and target arrays with a complex NumPy dtype.
+
+The code constructs a complex cosine target and restricts the search to `+`, `-`, `*`, and `cos`:
+
+```python
+import numpy as np
+
+from pysr import PySRRegressor
+
+X = np.random.randn(100, 1) + 1j * np.random.randn(100, 1)
+y = (1 + 2j) * np.cos(X[:, 0] * (0.5 - 0.2j))
+
+model = PySRRegressor(
+ binary_operators=["+", "-", "*"], unary_operators=["cos"], niterations=100,
+)
+
+model.fit(X, y)
+```
+
+The fitted constants belong to the complex-valued expression domain. Use `model.sympy()` to inspect the selected equation in SymPy form:
+
+```python
+model.sympy()
+```
+
+Pass an equation-row index to `predict`; `-1` selects the final row of `equations_`:
+
+```python
+model.predict(X, -1)
+```
+
+## Julia packages and types
+
+PySR delegates its symbolic-regression search to the pure Julia package [SymbolicRegression.jl](https://github.com/astroautomata/SymbolicRegression.jl). A custom operator can call functionality from another Julia package through that backend.
+
+Load the package that supplies the prime lookup before defining the operator. The target in this example is
+
+$$ y = p_{3x + 1} - 5, $$
+
+where $p_i$ denotes the $i$th prime and $x$ is the input feature. The lookup comes from [Primes.jl](https://github.com/JuliaMath/Primes.jl).
+
+First, import the handle to PySR's Julia runtime:
+
+```python
+from pysr import jl
+```
+
+`jl` is the juliacall runtime object, and `jl.seval` evaluates Julia source from Python. Use it to add `Primes.jl` to the active PySR environment:
+
+```python
+from pysr import jl
+
+jl.seval("""
+import Pkg
+Pkg.add("Primes")
+""")
+```
+
+Adding a package and importing its namespace are separate steps. Import `Primes` before defining `p`:
+
+```python
+from pysr import jl
+
+jl.seval("import Primes")
+```
+
+With the package available, define the unary operator used by the search:
+
+```python
+from pysr import jl
+
+jl.seval("""
+function p(i::T) where T
+ if (0.5 < i < 1000)
+ return T(Primes.prime(round(Int, i)))
+ else
+ return T(NaN)
+ end
+end
+""")
+```
+
+The operator rounds inputs within `0.5 < i < 1000` to prime indices and preserves the input type `T`. Outside that domain it returns `T(NaN)` to reject invalid evaluations.
+
+Call the Julia function through `jl.p` to build the prime values used by the synthetic data:
+
+```python
+from pysr import jl
+
+primes = {i: jl.p(i*1.0) for i in range(1, 999)}
+```
+
+Use the lookup table to generate noisy samples of the target:
+
+```python
+import numpy as np
+
+X = np.random.randint(0, 100, 100)[:, None]
+y = [primes[3*X[i, 0] + 1] - 5 + np.random.randn()*0.001 for i in range(100)]
+```
+
+Register `p` with the regular PySR operator interface. The `sympy_p` subclass and `extra_sympy_mappings` give SymPy a symbolic form for the custom function:
+
+```python
+from pysr import PySRRegressor
+import sympy
+
+class sympy_p(sympy.Function):
+ pass
+
+model = PySRRegressor(
+ binary_operators=["+", "-", "*", "/"],
+ unary_operators=["p"],
+ niterations=100,
+ extra_sympy_mappings={"p": sympy_p}
+)
+```
+
+Fit the model on the generated pairs:
+
+```python
+model.fit(X, y)
+```
+
+Since `p` rounds its argument, different real offsets can produce the same integer prime index, so the offset in a fitted expression need not equal `1` exactly. Inspect the selected equation in SymPy with:
+
+```python
+model.sympy()
+```
+
+## Custom value types
+
+`TypeSpec` declares a custom Julia value that can flow through an expression tree. Vectors, tensors, strings, structs, and other Julia value types can serve as payloads when their behavior is described by the type's hooks, operators, and loss.
+
+A model configured with `type_spec=TypeSpec(...)` needs three pieces:
+
+1. Define hooks that sample and change values. For continuously optimized constants, `scalar_constants` exposes scalar coordinates and `with_scalar_constants` rebuilds the value. A custom `mutate` hook handles changes that need a discrete or structural operation.
+2. Choose exactly one of `elementwise_loss`, `loss_function`, or `loss_function_expression`. The selected loss must accept the custom values and return a real-valued score. Elementwise-loss return types are inferred; a full objective requires an explicit `loss_type`.
+3. Supply operators through `operators={...}`, grouped by arity. Each operator must accept and return the declared value type so the type-stability check succeeds. Add an explicit return annotation such as `::Vec2` when Julia cannot infer it.
+
+The examples below combine these pieces for vector values, strings as discrete constants, and payloads with several shapes.
+
+### Vector-valued expression trees
+
+This example searches for an expression whose values are two-dimensional vectors:
+
+$$
+y = \operatorname{rotate90}(x_1) + 2x_2 +
+\begin{bmatrix}0.5 \\ -1.0\end{bmatrix}.
+$$
+
+The 128-by-2 DataFrame has two vector-valued features per row, and each entry of `y` is one target vector. The table dimensions describe samples and features, while each cell of the object array holds a single logical vector value.
+
+```python
+import numpy as np
+import pandas as pd
+
+from pysr import PySRRegressor, TypeSpec
+
+rng = np.random.default_rng(0)
+x1 = [rng.normal(size=2) for _ in range(128)]
+x2 = [rng.normal(size=2) for _ in range(128)]
+X = pd.DataFrame({"x1": x1, "x2": x2})
+y = np.empty(128, dtype=object)
+offset = np.array([0.5, -1.0])
+y[:] = [np.array([-a[1], a[0]]) + 2 * b + offset for a, b in zip(x1, x2)]
+```
+
+Declare `Vec2` with one `Vector{Float64}` field. `sample` creates a random payload of length two, while `scalar_constants` and `with_scalar_constants` expose and rebuild the entries used by continuous constant optimization:
+
+```python
+from pysr import TypeSpec
+
+type_spec = TypeSpec(
+ "Vec2",
+ fields={"data": "Vector{Float64}"},
+ sample="rng -> Vec2(randn(rng, 2))",
+ scalar_constants="value -> value.data",
+ with_scalar_constants=(
+ "(value, scalar_constants) -> Vec2(scalar_constants)"
+ ),
+)
+```
+
+PySR derives how values are initialized, mutated, validated, counted, packed, and unpacked from these hooks. Add `init` or `mutate` for specialized search behavior, or `string` for custom display, when the derived defaults are insufficient.
+
+`TypeSpec` generates the Julia value type from that declaration. Its equivalent Julia struct has one vector field:
+
+```julia
+struct Vec2
+ data::Vector{Float64}
+end
+```
+
+Define operators that consume and return `Vec2` values:
+
+```python
+operators = {
+ 1: [
+ "rotate90(a::Vec2) = Vec2([-a.data[2], a.data[1]])",
+ "double(a::Vec2) = Vec2(2a.data)",
+ ],
+ 2: ["add_vectors(a::Vec2, b::Vec2) = Vec2(a.data + b.data)"],
+}
+```
+
+The elementwise loss compares vector payloads by summing squared differences, which gives the real-valued score the search requires:
+
+```python
+from pysr import PySRRegressor
+
+model = PySRRegressor(
+ type_spec=type_spec,
+ operators=operators,
+ elementwise_loss="vector_loss(a::Vec2, b::Vec2) = sum(abs2, a.data - b.data)",
+ niterations=40,
+ populations=4,
+ maxsize=10,
+)
+
+model.fit(X, y)
+print(model.equations_)
+```
+
+The search space can represent the target as `add_vectors(add_vectors(rotate90(x1), double(x2)), Vec2([0.5, -1.0]))`. The vector constant carries both offset components, which the scalar hook pair exposes for joint optimization.
+
+
+Data generation code
+ +```python +import numpy as np + +rng = np.random.default_rng(20260817) +X = rng.uniform(-3.0, 3.0, size=(200, 5)) +y = 2.5382 * np.cos(X[:, 3]) + X[:, 0] ** 2 - 0.5 +``` + +
+
+
+String-valued expressions and discrete constants
+ +This example treats each string as a value and chooses a separator from a finite set during mutation. It lowercases the first input, uppercases the second, and joins them with the selected separator: + +```python +import numpy as np +import pandas as pd + +from pysr import PySRRegressor, TypeSpec + +X = pd.DataFrame( + { + "first": ["Py", "symbolic", "hello", "left"], + "second": ["SR", "regression", "world", "right"], + } +) +y = np.array( + [f"{a.lower()}-{b.upper()}" for a, b in X.itertuples(index=False)], + dtype=object, +) + +type_spec = TypeSpec( + "StringValue", + fields={"data": "String"}, + sample='rng -> StringValue(rand(rng, ("", "-", "_")))', + mutate=""" + mutate_string(rng, value, temperature) = StringValue(rand(rng, ("", "-", "_"))) + """, +) + +model = PySRRegressor( + type_spec=type_spec, + operators={ + 1: [ + "string_lowercase(x::StringValue) = StringValue(lowercase(x.data))", + "string_uppercase(x::StringValue) = StringValue(uppercase(x.data))", + ], + 2: [ + "string_concat(a::StringValue, b::StringValue) = StringValue(a.data * b.data)" + ], + }, + elementwise_loss=""" + string_loss(a::StringValue, b::StringValue) = Float64(Base.editdistance(a.data, b.data)) + """, + niterations=40, +) + +model.fit(X, y) +print(model.equations_) +``` + +The `sample` and `mutate` hooks choose discrete separator constants. Without a +`scalar_constants` and `with_scalar_constants` pair, BFGS leaves them unchanged. +Edit distance provides the real-valued loss for string outputs. + +
+
+
+The same custom-value interface can seed a search with `guesses`. In the vector example above,
+`Vec2` is the generated constructor from its `TypeSpec`, `add_vectors` is the declared binary
+operator, and `x1` is the first named feature column. A matching guess is
+`guesses=["add_vectors(x1, Vec2([1.0, 2.0]))"]`.
diff --git a/docs/src/migration.md b/docs/src/migration.md
new file mode 100644
index 000000000..3fa3b7437
--- /dev/null
+++ b/docs/src/migration.md
@@ -0,0 +1,207 @@
+# Migrating from v1 to v2
+
+PySR v2 upgrades the Julia backend from SymbolicRegression.jl v1.11 to v2.
+Most `PySRRegressor` constructor calls still work, but several search defaults
+changed. The old parametric-expression interface and v1 checkpoints do not
+carry over.
+
+Before upgrading, keep a copy of each run directory, including
+`hall_of_fame.csv`. Install v2 in a new Python environment if you need to keep a
+v1 environment available for old pickles or exact search reproduction.
+
+## Search defaults changed
+
+The following changes affect ordinary `PySRRegressor(...)` searches:
+
+| Setting | v1 default | v2 default |
+| --- | --- | --- |
+| `batching` | `False` | `"auto"` |
+| `batch_size` | `50` | `None` |
+| `annealing` | `False` | `True` |
+| `crossover_probability` | `0.0259` | `0.2` |
+| Mutation selection | Fixed v1 weights | New backend weights, adapted during the search |
+
+With `batching="auto"`, the backend enables batching for large datasets while
+still evaluating hall-of-fame candidates on the full dataset. With
+`batch_size=None`, it uses the full dataset for at most 1000 rows, then batches
+of 128, 256, or 512 as the dataset grows.
+
+The default mutation mix changed, including a new feature mutation that can
+replace the input referenced by a variable node. The
+`AdaptiveMutationWeightsPlugin` also updates mutation weights from observed
+success rates. A `weight_*` argument now supplies an initial weight, rather than
+a fixed weight, while that plugin is active.
+
+For a search close to the v1 defaults, set the changed values explicitly and
+omit the adaptive-mutation plugin:
+
+```python
+from pysr import AdaptiveParsimonyPlugin, PySRRegressor
+
+model = PySRRegressor(
+ batching=False,
+ batch_size=50,
+ annealing=False,
+ crossover_probability=0.0259,
+ default_plugins=[AdaptiveParsimonyPlugin()],
+ weight_add_node=2.47,
+ weight_insert_node=0.0112,
+ weight_delete_node=0.87,
+ weight_do_nothing=0.273,
+ weight_mutate_constant=0.0346,
+ weight_mutate_operator=0.293,
+ weight_mutate_feature=0.0,
+ weight_swap_operands=0.198,
+ weight_rotate_tree=4.26,
+ weight_randomize=0.000502,
+ weight_simplify=0.00209,
+ weight_optimize=0.0,
+ weight_backsolve=0.0,
+)
+```
+
+This preserves the v1 search configuration, but it does not make a run
+bit-for-bit identical. Backend implementation changes still alter random-number
+consumption and search trajectories.
+
+Two other backend changes can alter results:
+
+- Simplification now recomputes an expression's cost from scratch. Equivalent
+ expressions can therefore appear in a different hall-of-fame order.
+- Constant optimization can restart zero-valued constants, allowing trajectories
+ that v1 could not reach.
+
+Loss functions are now validated when `fit` is called. An `elementwise_loss` or
+`loss_function` with the wrong number of arguments raises a `ValueError` before
+the search starts.
+
+## Parametric expressions
+
+`ParametricExpressionSpec` and the `category=` argument to `fit` were removed.
+Use `TemplateExpressionSpec.parameters`, append the category to `X`, and index
+parameters from the template.
+
+A v1 model such as:
+
+```python
+from pysr import ParametricExpressionSpec, PySRRegressor
+
+model = PySRRegressor(
+ expression_spec=ParametricExpressionSpec(max_parameters=2),
+)
+model.fit(X, y, category=category)
+```
+
+becomes:
+
+```python
+import numpy as np
+from pysr import PySRRegressor, TemplateExpressionSpec
+
+n_categories = len(np.unique(category))
+spec = TemplateExpressionSpec(
+ combine="f(x1, x2, p1[category], p2[category])",
+ expressions=["f"],
+ variable_names=["x1", "x2", "category"],
+ parameters={"p1": n_categories, "p2": n_categories},
+)
+
+# Julia indexing starts at 1. Add one if Python categories start at 0.
+X_with_category = np.column_stack([X, category + 1])
+model = PySRRegressor(expression_spec=spec)
+model.fit(X_with_category, y)
+```
+
+The category column must also be present when calling `predict`.
+
+## Template expressions
+
+The deprecated `function_symbols=...` and positional constructor forms of
+`TemplateExpressionSpec` were replaced by explicit template fields. State the
+combining expression, inner expressions, and variable names directly:
+
+```python
+spec = TemplateExpressionSpec(
+ combine="sin(f(x1, x2)) + g(x3)^2",
+ expressions=["f", "g"],
+ variable_names=["x1", "x2", "x3"],
+)
+```
+
+This form also supports optimizable parameter vectors through
+`parameters={"p": length}`. See [Template expressions](examples.md#template-expressions)
+for complete examples.
+
+## Saved models and checkpoints
+
+v2 uses checkpoint schema 3. A v1 pickle has no compatible schema marker, so
+`pickle.load(...)` and `PySRRegressor.from_file(...)` reject it with an
+`Unsupported PySR checkpoint schema` error. Warm starts cannot cross the v1 to
+v2 boundary. Re-run the search under v2 if you need to continue evolution.
+
+For standard expressions, you can reconstruct the discovered equations from a
+copy of `hall_of_fame.csv`. Put the CSV in a run directory without the
+incompatible `checkpoint.pkl`, then provide the original operators and feature
+metadata:
+
+```python
+model = PySRRegressor.from_file(
+ run_directory="v1_csv_copy",
+ binary_operators=["+", "-", "*", "/"],
+ unary_operators=["sin", "cos"],
+ n_features_in=3,
+ feature_names_in=["x0", "x1", "x2"],
+)
+```
+
+Custom expression specifications may require a fresh search because their Julia
+objects cannot be reconstructed from the CSV alone.
+
+## v1 spellings that still work
+
+These forms emit a deprecation warning but retain their old behavior:
+
+- camelCase keyword arguments such as `fractionReplaced`, `npop`,
+ `weightMutateConstant`, and `crossoverProbability`; use snake_case instead.
+- `loss=` and `full_objective=`; use `elementwise_loss=` and `loss_function=`.
+- `multithreading=True` or `False`; use `parallelism="multithreading"` or
+ `parallelism="serial"`. Use `parallelism="multiprocessing", procs=N` for
+ process workers.
+- `.equations` and `.raw_julia_state_`; use `.equations_` and `.julia_state_`.
+- The `pysr()` function, `pysr.install()`, and `pysr.init_julia()`; use
+ `PySRRegressor`. Dependencies are installed when PySR is imported.
+
+These constructor arguments are accepted but ignored with a warning:
+
+- `julia_project` and `julia_kwargs`. Configure the backend through `juliapkg`;
+ see [Customizing the backend](backend.md).
+- `weights`, `variable_names`, and `Xresampled`. Pass them to `fit` instead.
+
+## New APIs in v2
+
+- **N-ary operators:** use
+ `operators={1: ["sin"], 2: ["+", "*"], 3: ["clamp"]}`. The legacy
+ `binary_operators` and `unary_operators` arguments still work, but cannot be
+ combined with `operators`. Constraints for n-ary operators use n-tuples, such
+ as `{"clamp": (-1, -1, 1)}`.
+- **First-class mutations, crossovers, and plugins:** configure weighted
+ instances with `mutations=` and `crossovers=`, and compose search behavior
+ with `plugins=`. See the [API reference](api.md#mutations).
+- **Custom value types:** `TypeSpec` supports expressions over strings, vectors,
+ tensors, structs, and other Julia value types. See
+ [Custom value types](examples.md#custom-value-types).
+- **Template parameters:** `TemplateExpressionSpec.parameters` supports learned
+ scalar or category-indexed parameter vectors.
+- **Equation guesses:** seed a search with `guesses=["sin(x1 * 2.1)", ...]`.
+ `fraction_replaced_guesses` controls their injection rate.
+- **Autodiff backends:** `autodiff_backend="Enzyme"` is available as an
+ experimental alternative to `"Zygote"`.
+- **Worker controls:** `worker_imports` loads Julia packages on multiprocessing
+ workers, and `worker_timeout` controls worker restart timeouts.
+- **Additional mutations:** `weight_mutate_feature` controls feature mutation.
+ Backsolve is disabled by default and can be configured with
+ `BacksolveMutation(...)` through `mutations=`.
+
+## Environment changes
+
+- The minimum versions are Python 3.10 and Julia 1.10.
diff --git a/docs/src/options.md b/docs/src/options.md
index 7a8f5c11e..3590808bc 100644
--- a/docs/src/options.md
+++ b/docs/src/options.md
@@ -245,7 +245,7 @@ The default loss is mean-square error, and weighted mean-square error.
One can pass an arbitrary Julia string to define a custom loss, using,
e.g., `elementwise_loss="myloss(x, y) = abs(x - y)^1.5"`. For more details,
see the
-[Losses](https://milescranmer.github.io/SymbolicRegression.jl/dev/losses/)
+[Losses](https://julia.pysr.ai/dev/losses)
page for SymbolicRegression.jl.
Here are some additional examples:
@@ -276,7 +276,7 @@ model = PySRRegressor(..., weights=weights, elementwise_loss="myloss(x, y, w) =
model.fit(..., weights=weights)
```
-Built-in loss (faster) (see [losses](https://ai.damtp.cam.ac.uk/symbolicregression/dev/losses/)).
+Built-in loss (faster) (see [losses](https://julia.pysr.ai/dev/losses)).
This one computes the L3 norm:
```python
diff --git a/docs/src/public/apple-touch-icon.png b/docs/src/public/apple-touch-icon.png
new file mode 100644
index 000000000..439235373
Binary files /dev/null and b/docs/src/public/apple-touch-icon.png differ
diff --git a/docs/src/public/favicon-16x16.png b/docs/src/public/favicon-16x16.png
new file mode 100644
index 000000000..d6a26a31d
Binary files /dev/null and b/docs/src/public/favicon-16x16.png differ
diff --git a/docs/src/public/favicon-32x32.png b/docs/src/public/favicon-32x32.png
new file mode 100644
index 000000000..d46b06535
Binary files /dev/null and b/docs/src/public/favicon-32x32.png differ
diff --git a/docs/src/public/favicon.ico b/docs/src/public/favicon.ico
new file mode 100644
index 000000000..ba1cdba34
Binary files /dev/null and b/docs/src/public/favicon.ico differ
diff --git a/docs/src/public/favicon.png b/docs/src/public/favicon.png
index 5fcc0ab9f..69498a477 100644
Binary files a/docs/src/public/favicon.png and b/docs/src/public/favicon.png differ
diff --git a/docs/src/slurm.md b/docs/src/slurm.md
new file mode 100644
index 000000000..ec1c395bf
--- /dev/null
+++ b/docs/src/slurm.md
@@ -0,0 +1,50 @@
+# Slurm (Multi-node)
+
+PySR supports running across multiple nodes in an existing Slurm allocation with
+`cluster_manager="slurm"`. Request the resources with `sbatch` or `salloc`, then
+PySR will launch Julia workers across the allocation using
+`SlurmClusterManager.jl`.
+
+For example, this job requests three workers on each of two nodes:
+
+```bash
+#!/bin/bash
+#SBATCH --job-name=pysr
+#SBATCH --partition=normal
+#SBATCH --nodes=2
+#SBATCH --ntasks-per-node=3
+#SBATCH --time=01:00:00
+
+set -euo pipefail
+python pysr_script.py
+```
+
+The corresponding `pysr_script.py` is:
+
+```python
+import numpy as np
+from pysr import PySRRegressor
+
+X = np.random.RandomState(0).randn(1000, 2)
+y = X[:, 0] + 2 * X[:, 1]
+
+model = PySRRegressor(
+ niterations=200,
+ populations=2,
+ parallelism="multiprocessing",
+ cluster_manager="slurm",
+ procs=6,
+)
+model.fit(X, y)
+print(model)
+```
+
+Submit the job with:
+
+```bash
+sbatch pysr_job.sh
+```
+
+`procs` must equal the allocation's total task count, given by `--ntasks` or
+`--nodes` multiplied by `--ntasks-per-node`. Run the Python script once inside
+the allocation; do not wrap it in `srun`.
diff --git a/docs/src/tuning.md b/docs/src/tuning.md
index 443789e49..ec995d29b 100644
--- a/docs/src/tuning.md
+++ b/docs/src/tuning.md
@@ -6,11 +6,11 @@ First, my general tips would be to avoid using redundant operators, like how `po
When running PySR, I usually do the following:
-I run from IPython (Jupyter Notebooks don't work as well[^1]) on the head node of a slurm cluster. Passing `cluster_manager="slurm"` will make PySR set up a run over the entire allocation. I set `procs` equal to the total number of cores over my entire allocation.
+I run from IPython (Jupyter Notebooks don't work as well[^1]) on the head node of a slurm cluster. Passing `cluster_manager="slurm"` will make PySR set up a run over the entire allocation. I set `procs` equal to the total number of tasks across my entire allocation (see the [Slurm page](slurm.md) for a complete multi-node example).
-I use the [tensorboard feature](https://ai.damtp.cam.ac.uk/pysr/examples/#12-using-tensorboard-for-logging) for experiment tracking.
+I use the [tensorboard feature](/examples/instrumentation#using-tensorboard-for-logging) for experiment tracking.
-[^1]: Jupyter Notebooks are supported by PySR, but miss out on some useful features available in IPython and Python: the progress bar, and early stopping with "q". In Jupyter you cannot interrupt a search once it has started; you have to restart the kernel. See [this issue](https://github.com/astroautomata/PySR/issues/260) for updates.
+[^1]: Jupyter Notebooks are supported by PySR, but miss out on some useful features available in IPython and Python: the progress bar, and early stopping with "q". Interrupting a search in Jupyter returns the equations found so far.
1. I start by using the default parameters.
2. I use only the operators I think it needs and no more.
diff --git a/docs/test_deploy_versions.jl b/docs/test_deploy_versions.jl
new file mode 100644
index 000000000..886bd939d
--- /dev/null
+++ b/docs/test_deploy_versions.jl
@@ -0,0 +1,163 @@
+using Test
+
+include(joinpath(@__DIR__, "deploy_versions.jl"))
+
+@testset "stable ownership" begin
+ tags = [v"2.1.0", v"2.0.0", v"1.5.10"]
+
+ @test claims_stable("v2.1.0"; all_tagged_versions = tags, log = false)
+ @test !claims_stable("v2.0.0"; all_tagged_versions = tags, log = false)
+ # determine_bases would report this alpha as owning v2.0.0
+ @test !claims_stable("v2.0.0a2"; all_tagged_versions = tags, log = false)
+ @test !claims_stable("dev"; all_tagged_versions = tags, log = false)
+ @test !claims_stable("previews/PR123"; all_tagged_versions = tags, log = false)
+end
+
+@testset "stable deployment root" begin
+ @test stable_deploy_root("/tmp/gh-pages/stable") == "/tmp/gh-pages"
+ @test stable_deploy_root("/tmp/gh-pages/stable/") == "/tmp/gh-pages"
+end
+
+@testset "stable redirects" begin
+ pages = ["index", "examples", "404", "api"]
+ files = stable_redirect_files(pages, "v2.2.1")
+
+ @test files["api.html"] == """
+
+"""
+ @test files["api/index.html"] == """
+
+"""
+ @test files["stable/api.html"] == """
+
+"""
+ @test files["stable/index.html"] == """
+
+"""
+
+ @test sort(collect(keys(files))) == [
+ "api.html",
+ "api/index.html",
+ "examples.html",
+ "examples/index.html",
+ "stable/api.html",
+ "stable/examples.html",
+ "stable/index.html",
+ ]
+ @test collect(files) == collect(stable_redirect_files(reverse(pages), "v2.2.1"))
+ root_stubs = filter(((path, _),) -> !startswith(path, "stable/"), files)
+ @test all(!occursin(r"v\d", content) for content in values(root_stubs))
+end
+
+@testset "stale redirects" begin
+ files = stable_redirect_files(["api"], "v2.2.1")
+ root_stubs = filter(((path, _),) -> !startswith(path, "stable/"), files)
+
+ root = mktempdir()
+ for (path, content) in root_stubs
+ mkpath(joinpath(root, Base.dirname(path)))
+ write(joinpath(root, path), content)
+ end
+ # A page dropped since the previous release, in both of its stub forms.
+ for (path, content) in stable_redirect_files(["gone"], "v2.2.0")
+ startswith(path, "stable/") && continue
+ mkpath(joinpath(root, Base.dirname(path)))
+ write(joinpath(root, path), content)
+ end
+ # Everything a deployment leaves at the root besides our stubs.
+ write(joinpath(root, "index.html"), "\n")
+ write(joinpath(root, "versions.js"), "var DOC_VERSIONS = [];\n")
+ write(joinpath(root, "CNAME"), "pysr.ai\n")
+ for version in ("v2.2.1", "dev", "stable")
+ mkpath(joinpath(root, version))
+ write(joinpath(root, version, "index.html"), "\n")
+ end
+ write(joinpath(root, "stable", "index.html"), files["stable/index.html"])
+
+ @test stale_root_stubs(root, keys(root_stubs)) == ["gone.html", "gone/index.html"]
+ @test read(joinpath(root, "CNAME"), String) == "pysr.ai\n"
+ @test isempty(stale_root_stubs(root, keys(stable_redirect_files(["api", "gone"], "v2.2.1"))))
+end
+
+@testset "nested redirects" begin
+ files = stable_redirect_files(["api", "examples/objectives", "guide/index"], "v2.3.0")
+
+ @test files["examples/objectives.html"] == """$(STUB_MARKER)
+
+"""
+ @test files["examples/objectives/index.html"] == """$(STUB_MARKER)
+
+"""
+ @test files["stable/examples/objectives.html"] == """$(STUB_MARKER)
+
+"""
+ # A page that is already a directory index gets no `