1725881820

Build Type

github

Committed by GitHub

Commit Message

<a href="https://github.com/MilesCranmer/PySR/commit/<a class=hub.com/MilesCranmer/PySR/commit/<a class="double-link" href="https://git"><a class=hub.com/MilesCranmer/PySR/commit/<a class="double-link" href="https://git"><a class=hub.com/MilesCranmer/PySR/commit/<a class="double-link" href="https://git"><a class=hub.com/MilesCranmer/PySR/commit/<a class="double-link" href="https://git"><a class=hub.com/MilesCranmer/PySR/commit/ae7c60a8b9c2cf1e5508b7ddf35a74b57b8a9fc5">ae7c60a8b">&lt;a href=&quot;https://github.com/MilesCranmer/PySR/commit/</a><a class="double-link" href="https://github.com/MilesCranmer/PySR/commit/&lt;a class=&quot;double-link&quot; href=&quot;https://git">&lt;a class=</a>hub.com/MilesCranmer/PySR/commit/&lt;a class=&quot;double-link&quot; href=&quot;https://git">&lt;a class=</a>hub.com/MilesCranmer/PySR/commit/&lt;a class=&quot;double-link&quot; href=&quot;https://git">&lt;a class=</a>hub.com/MilesCranmer/PySR/commit/ae7c60a8b9c2cf1e5508b7ddf35a74b57b8a9fc5">ae7c60a8b</a><a href="https://github.com/MilesCranmer/PySR/commit/ae7c60a8b9c2cf1e5508b7ddf35a74b57b8a9fc5">&lt;a href=&quot;https://github.com/MilesCranmer/PySR/commit/ae7c60a8b9c2cf1e5508b7ddf35a74b57b8a9fc5&quot;&gt;&amp;quot;&amp;gt;&amp;amp;lt;a href=&amp;amp;quot;https://github.com/MilesCranmer/PySR/commit/&amp;lt;/a&amp;gt;&amp;lt;a class=&amp;quot;double-link&amp;quot; href=&amp;quot;https://github.com/MilesCranmer/PySR/commit/&amp;amp;lt;a class=&amp;amp;quot;double-link&amp;amp;quot; href=&amp;amp;quot;https://git&amp;quot;&amp;gt;&amp;amp;lt;a class=&amp;lt;/a&amp;gt;hub.com/MilesCranmer/PySR/commit/ae7c60a8b9c2cf1e5508b7ddf35a74b57b8a9fc5&amp;quot;&amp;gt;ae7c60a8b&amp;lt;/a&amp;gt;&amp;lt;a href=&amp;quot;https://github.com/MilesCranmer/PySR/commit/ae7c60a8b9c2cf1e5508b7ddf35a74b57b8a9fc5&amp;quot;&amp;gt;&amp;amp;lt;a href=&amp;amp;quot;https://github.com/MilesCranmer/PySR/commit/ae7c60a8b9c2cf1e5508b7ddf35a74b5&amp;lt;/a&amp;gt;7b8a9</a>fc5&quot;&gt;&amp;quot;&amp;gt;Merge &amp;lt;/a&amp;gt;&amp;lt;a class=&amp;quot;double-link&amp;quot; href=&amp;quot;https://github.com/MilesCranmer/PySR/commit/&amp;lt;a class=&amp;quot;double-link&amp;quot; href=&amp;quot;https://github.com/MilesCranmer/PySR/commit/&lt;a class=&quot;double-link&quot; href=&quot;https://github.com/MilesCranmer/PySR/commit/ceeb4200e&quot;&gt;ceeb420... (continued)

Pull Request Pull Request #87: Switch Julia interface to PyJulia

Run Details

79 of 103 new or added lines in 2 files covered. (76.7%)

455 of 526 relevant lines covered (86.5%)

7.79 hits per line

Source File
Press 'n' to go to next uncovered line, 'b' for previous

84.28

/pysr/sr.py

1	import os	9✔
2	import sys	9✔
3	import numpy as np	9✔
4	import pandas as pd	9✔
5	import sympy	9✔
6	from sympy import sympify, lambdify	9✔
7	import subprocess	9✔
8	import tempfile	9✔
9	import shutil	9✔
10	from pathlib import Path	9✔
11	from datetime import datetime	9✔
12	import warnings	9✔
13	from multiprocessing import cpu_count	9✔
14
15	is_julia_warning_silenced = False	9✔
16
17
18	def install(julia_project=None):	9✔
NEW 19	import julia	×
20
NEW 21	julia.install()	×
22
NEW 23	julia_project = _get_julia_project(julia_project)	×
24
NEW 25	init_julia()	×
NEW 26	from julia import Pkg	×
27
NEW 28	Pkg.activate(f"{_escape_filename(julia_project)}")	×
NEW 29	Pkg.update()	×
NEW 30	Pkg.instantiate()	×
NEW 31	Pkg.precompile()	×
32
33
34	Main = None	9✔
35	global_state = dict(	9✔
36	equation_file="hall_of_fame.csv",
37	n_features=None,
38	variable_names=[],
39	extra_sympy_mappings={},
40	extra_torch_mappings={},
41	extra_jax_mappings={},
42	output_jax_format=False,
43	output_torch_format=False,
44	multioutput=False,
45	nout=1,
46	selection=None,
47	raw_julia_output=None,
48	)
49
50	already_ran = False	9✔
51
52	sympy_mappings = {	9✔
53	"div": lambda x, y: x / y,
54	"mult": lambda x, y: x * y,
55	"sqrt_abs": lambda x: sympy.sqrt(abs(x)),
56	"square": lambda x: x ** 2,
57	"cube": lambda x: x ** 3,
58	"plus": lambda x, y: x + y,
59	"sub": lambda x, y: x - y,
60	"neg": lambda x: -x,
61	"pow": lambda x, y: abs(x) ** y,
62	"cos": sympy.cos,
63	"sin": sympy.sin,
64	"tan": sympy.tan,
65	"cosh": sympy.cosh,
66	"sinh": sympy.sinh,
67	"tanh": sympy.tanh,
68	"exp": sympy.exp,
69	"acos": sympy.acos,
70	"asin": sympy.asin,
71	"atan": sympy.atan,
72	"acosh": lambda x: sympy.acosh(abs(x) + 1),
73	"acosh_abs": lambda x: sympy.acosh(abs(x) + 1),
74	"asinh": sympy.asinh,
75	"atanh": lambda x: sympy.atanh(sympy.Mod(x + 1, 2) - 1),
76	"atanh_clip": lambda x: sympy.atanh(sympy.Mod(x + 1, 2) - 1),
77	"abs": abs,
78	"mod": sympy.Mod,
79	"erf": sympy.erf,
80	"erfc": sympy.erfc,
81	"log_abs": lambda x: sympy.log(abs(x)),
82	"log10_abs": lambda x: sympy.log(abs(x), 10),
83	"log2_abs": lambda x: sympy.log(abs(x), 2),
84	"log1p_abs": lambda x: sympy.log(abs(x) + 1),
85	"floor": sympy.floor,
86	"ceil": sympy.ceiling,
87	"sign": sympy.sign,
88	"gamma": sympy.gamma,
89	}
90
91
92	def pysr(	9✔
93	X,
94	y,
95	weights=None,
96	binary_operators=None,
97	unary_operators=None,
98	procs=cpu_count(),
99	loss="L2DistLoss()",
100	populations=20,
101	niterations=100,
102	ncyclesperiteration=300,
103	alpha=0.1,
104	annealing=False,
105	fractionReplaced=0.10,
106	fractionReplacedHof=0.10,
107	npop=1000,
108	parsimony=1e-4,
109	migration=True,
110	hofMigration=True,
111	shouldOptimizeConstants=True,
112	topn=10,
113	weightAddNode=1,
114	weightInsertNode=3,
115	weightDeleteNode=3,
116	weightDoNothing=1,
117	weightMutateConstant=10,
118	weightMutateOperator=1,
119	weightRandomize=1,
120	weightSimplify=0.01,
121	perturbationFactor=1.0,
122	extra_sympy_mappings=None,
123	extra_torch_mappings=None,
124	extra_jax_mappings=None,
125	equation_file=None,
126	verbosity=1e9,
127	progress=None,
128	maxsize=20,
129	fast_cycle=False,
130	maxdepth=None,
131	variable_names=None,
132	batching=False,
133	batchSize=50,
134	select_k_features=None,
135	warmupMaxsizeBy=0.0,
136	constraints=None,
137	useFrequency=True,
138	tempdir=None,
139	delete_tempfiles=True,
140	julia_project=None,
141	update=True,
142	temp_equation_file=False,
143	output_jax_format=False,
144	output_torch_format=False,
145	optimizer_algorithm="BFGS",
146	optimizer_nrestarts=3,
147	optimize_probability=1.0,
148	optimizer_iterations=10,
149	tournament_selection_n=10,
150	tournament_selection_p=1.0,
151	denoise=False,
152	Xresampled=None,
153	precision=32,
154	multithreading=None,
155	**kwargs,
156	):
157	"""Run symbolic regression to fit f(X[i, :]) ~ y[i] for all i.
158	Note: most default parameters have been tuned over several example
159	equations, but you should adjust `niterations`,
160	`binary_operators`, `unary_operators` to your requirements.
161	You can view more detailed explanations of the options on the
162	[options page](https://pysr.readthedocs.io/en/latest/docs/options/) of the documentation.
163
164	:param X: 2D array. Rows are examples, columns are features. If pandas DataFrame, the columns are used for variable names (so make sure they don't contain spaces).
165	:type X: np.ndarray/pandas.DataFrame
166	:param y: 1D array (rows are examples) or 2D array (rows are examples, columns are outputs). Putting in a 2D array will trigger a search for equations for each feature of y.
167	:type y: np.ndarray
168	:param weights: same shape as y. Each element is how to weight the mean-square-error loss for that particular element of y.
169	:type weights: np.ndarray
170	:param binary_operators: List of strings giving the binary operators in Julia's Base. Default is ["+", "-", "*", "/",].
171	:type binary_operators: list
172	:param unary_operators: Same but for operators taking a single scalar. Default is [].
173	:type unary_operators: list
174	:param procs: Number of processes (=number of populations running).
175	:type procs: int
176	:param loss: String of Julia code specifying the loss function. Can either be a loss from LossFunctions.jl, or your own loss written as a function. Examples of custom written losses include: `myloss(x, y) = abs(x-y)` for non-weighted, or `myloss(x, y, w) = w*abs(x-y)` for weighted. Among the included losses, these are as follows. Regression: `LPDistLoss{P}()`, `L1DistLoss()`, `L2DistLoss()` (mean square), `LogitDistLoss()`, `HuberLoss(d)`, `L1EpsilonInsLoss(ϵ)`, `L2EpsilonInsLoss(ϵ)`, `PeriodicLoss(c)`, `QuantileLoss(τ)`. Classification: `ZeroOneLoss()`, `PerceptronLoss()`, `L1HingeLoss()`, `SmoothedL1HingeLoss(γ)`, `ModifiedHuberLoss()`, `L2MarginLoss()`, `ExpLoss()`, `SigmoidLoss()`, `DWDMarginLoss(q)`.
177	:type loss: str
178	:param populations: Number of populations running.
179	:type populations: int
180	:param niterations: Number of iterations of the algorithm to run. The best equations are printed, and migrate between populations, at the end of each.
181	:type niterations: int
182	:param ncyclesperiteration: Number of total mutations to run, per 10 samples of the population, per iteration.
183	:type ncyclesperiteration: int
184	:param alpha: Initial temperature.
185	:type alpha: float
186	:param annealing: Whether to use annealing. You should (and it is default).
187	:type annealing: bool
188	:param fractionReplaced: How much of population to replace with migrating equations from other populations.
189	:type fractionReplaced: float
190	:param fractionReplacedHof: How much of population to replace with migrating equations from hall of fame.
191	:type fractionReplacedHof: float
192	:param npop: Number of individuals in each population
193	:type npop: int
194	:param parsimony: Multiplicative factor for how much to punish complexity.
195	:type parsimony: float
196	:param migration: Whether to migrate.
197	:type migration: bool
198	:param hofMigration: Whether to have the hall of fame migrate.
199	:type hofMigration: bool
200	:param shouldOptimizeConstants: Whether to numerically optimize constants (Nelder-Mead/Newton) at the end of each iteration.
201	:type shouldOptimizeConstants: bool
202	:param topn: How many top individuals migrate from each population.
203	:type topn: int
204	:param perturbationFactor: Constants are perturbed by a max factor of (perturbationFactor*T + 1). Either multiplied by this or divided by this.
205	:type perturbationFactor: float
206	:param weightAddNode: Relative likelihood for mutation to add a node
207	:type weightAddNode: float
208	:param weightInsertNode: Relative likelihood for mutation to insert a node
209	:type weightInsertNode: float
210	:param weightDeleteNode: Relative likelihood for mutation to delete a node
211	:type weightDeleteNode: float
212	:param weightDoNothing: Relative likelihood for mutation to leave the individual
213	:type weightDoNothing: float
214	:param weightMutateConstant: Relative likelihood for mutation to change the constant slightly in a random direction.
215	:type weightMutateConstant: float
216	:param weightMutateOperator: Relative likelihood for mutation to swap an operator.
217	:type weightMutateOperator: float
218	:param weightRandomize: Relative likelihood for mutation to completely delete and then randomly generate the equation
219	:type weightRandomize: float
220	:param weightSimplify: Relative likelihood for mutation to simplify constant parts by evaluation
221	:type weightSimplify: float
222	:param equation_file: Where to save the files (.csv separated by \|)
223	:type equation_file: str
224	:param verbosity: What verbosity level to use. 0 means minimal print statements.
225	:type verbosity: int
226	:param progress: Whether to use a progress bar instead of printing to stdout.
227	:type progress: bool
228	:param maxsize: Max size of an equation.
229	:type maxsize: int
230	:param maxdepth: Max depth of an equation. You can use both maxsize and maxdepth. maxdepth is by default set to = maxsize, which means that it is redundant.
231	:type maxdepth: int
232	:param fast_cycle: (experimental) - batch over population subsamples. This is a slightly different algorithm than regularized evolution, but does cycles 15% faster. May be algorithmically less efficient.
233	:type fast_cycle: bool
234	:param variable_names: a list of names for the variables, other than "x0", "x1", etc.
235	:type variable_names: list
236	:param batching: whether to compare population members on small batches during evolution. Still uses full dataset for comparing against hall of fame.
237	:type batching: bool
238	:param batchSize: the amount of data to use if doing batching.
239	:type batchSize: int
240	:param select_k_features: whether to run feature selection in Python using random forests, before passing to the symbolic regression code. None means no feature selection; an int means select that many features.
241	:type select_k_features: None/int
242	:param warmupMaxsizeBy: whether to slowly increase max size from a small number up to the maxsize (if greater than 0). If greater than 0, says the fraction of training time at which the current maxsize will reach the user-passed maxsize.
243	:type warmupMaxsizeBy: float
244	:param constraints: dictionary of int (unary) or 2-tuples (binary), this enforces maxsize constraints on the individual arguments of operators. E.g., `'pow': (-1, 1)` says that power laws can have any complexity left argument, but only 1 complexity exponent. Use this to force more interpretable solutions.
245	:type constraints: dict
246	:param useFrequency: whether to measure the frequency of complexities, and use that instead of parsimony to explore equation space. Will naturally find equations of all complexities.
247	:type useFrequency: bool
248	:param tempdir: directory for the temporary files
249	:type tempdir: str/None
250	:param delete_tempfiles: whether to delete the temporary files after finishing
251	:type delete_tempfiles: bool
252	:param julia_project: a Julia environment location containing a Project.toml (and potentially the source code for SymbolicRegression.jl). Default gives the Python package directory, where a Project.toml file should be present from the install.
253	:type julia_project: str/None
254	:param update: Whether to automatically update Julia packages.
255	:type update: bool
256	:param temp_equation_file: Whether to put the hall of fame file in the temp directory. Deletion is then controlled with the delete_tempfiles argument.
257	:type temp_equation_file: bool
258	:param output_jax_format: Whether to create a 'jax_format' column in the output, containing jax-callable functions and the default parameters in a jax array.
259	:type output_jax_format: bool
260	:param output_torch_format: Whether to create a 'torch_format' column in the output, containing a torch module with trainable parameters.
261	:type output_torch_format: bool
262	:param tournament_selection_n: Number of expressions to consider in each tournament.
263	:type tournament_selection_n: int
264	:param tournament_selection_p: Probability of selecting the best expression in each tournament. The probability will decay as p*(1-p)^n for other expressions, sorted by loss.
265	:type tournament_selection_p: float
266	:param denoise: Whether to use a Gaussian Process to denoise the data before inputting to PySR. Can help PySR fit noisy data.
267	:type denoise: bool
268	:param precision: What precision to use for the data. By default this is 32 (float32), but you can select 64 or 16 as well.
269	:type precision: int
270	:param multithreading: Use multithreading instead of distributed backend. Default is yes. Using procs=0 will turn off both.
271	:type multithreading: bool
272	:param **kwargs: Other options passed to SymbolicRegression.Options, for example, if you modify SymbolicRegression.jl to include additional arguments.
273	:type **kwargs: dict
274	:returns: Results dataframe, giving complexity, MSE, and equations (as strings), as well as functional forms. If list, each element corresponds to a dataframe of equations for each output.
275	:type: pd.DataFrame/list
276	"""
277	global already_ran
278
279	if binary_operators is None:	9✔
280	binary_operators = "+ * - /".split(" ")	9✔
281	if unary_operators is None:	9✔
282	unary_operators = []	9✔
283	if extra_sympy_mappings is None:	9✔
284	extra_sympy_mappings = {}	9✔
285	if variable_names is None:	9✔
286	variable_names = []	9✔
287	if constraints is None:	9✔
288	constraints = {}	9✔
289	if multithreading is None:	9✔
290	# Default is multithreading=True, unless explicitly set,
291	# or procs is set to 0 (serial mode).
292	multithreading = procs != 0	9✔
293
294	global Main
295	if Main is None:	9✔
296	if multithreading:	9✔
297	os.environ["JULIA_NUM_THREADS"] = str(procs)	9✔
298
299	Main = init_julia()	9✔
300
301	buffer_available = "buffer" in sys.stdout.__dir__()	9✔
302
303	if progress is not None:	9✔
304	if progress and not buffer_available:	×
305	warnings.warn(	×
306	"Note: it looks like you are running in Jupyter. The progress bar will be turned off."
307	)
308	progress = False	×
309	else:
310	progress = buffer_available	9✔
311
312	assert optimizer_algorithm in ["NelderMead", "BFGS"]	9✔
313	assert tournament_selection_n < npop	9✔
314
315	if isinstance(X, pd.DataFrame):	9✔
316	variable_names = list(X.columns)	×
317	X = np.array(X)	×
318
319	if len(X.shape) == 1:	9✔
320	X = X[:, None]	×
321
322	if len(variable_names) == 0:	9✔
323	variable_names = [f"x{i}" for i in range(X.shape[1])]	9✔
324
325	if extra_jax_mappings is not None:	9✔
326	for value in extra_jax_mappings.values():	×
327	if not isinstance(value, str):	×
328	raise NotImplementedError(	×
329	"extra_jax_mappings must have keys that are strings! e.g., {sympy.sqrt: 'jnp.sqrt'}."
330	)
331
332	if extra_torch_mappings is not None:	9✔
333	for value in extra_jax_mappings.values():	×
334	if not callable(value):	×
335	raise NotImplementedError(	×
336	"extra_torch_mappings must be callable functions! e.g., {sympy.sqrt: torch.sqrt}."
337	)
338
339	use_custom_variable_names = len(variable_names) != 0	9✔
340
341	_check_assertions(	9✔
342	X,
343	binary_operators,
344	unary_operators,
345	use_custom_variable_names,
346	variable_names,
347	weights,
348	y,
349	)
350
351	if len(X) > 10000 and not batching:	9✔
352	warnings.warn(	×
353	"Note: you are running with more than 10,000 datapoints. You should consider turning on batching (https://pysr.readthedocs.io/en/latest/docs/options/#batching). You should also reconsider if you need that many datapoints. Unless you have a large amount of noise (in which case you should smooth your dataset first), generally < 10,000 datapoints is enough to find a functional form with symbolic regression. More datapoints will lower the search speed."
354	)
355
356	if maxsize > 40:	9✔
357	warnings.warn(	×
358	"Note: Using a large maxsize for the equation search will be exponentially slower and use significant memory. You should consider turning `useFrequency` to False, and perhaps use `warmupMaxsizeBy`."
359	)
360	if maxsize < 7:	9✔
361	raise NotImplementedError("PySR requires a maxsize of at least 7")	×
362
363	X, variable_names, selection = _handle_feature_selection(	9✔
364	X, select_k_features, use_custom_variable_names, variable_names, y
365	)
366
367	if maxdepth is None:	9✔
368	maxdepth = maxsize	9✔
369	if isinstance(binary_operators, str):	9✔
370	binary_operators = [binary_operators]	×
371	if isinstance(unary_operators, str):	9✔
372	unary_operators = [unary_operators]	×
373
374	if len(y.shape) == 1 or (len(y.shape) == 2 and y.shape[1] == 1):	9✔
375	multioutput = False	9✔
376	nout = 1	9✔
377	y = y.reshape(-1)	9✔
378	elif len(y.shape) == 2:	9✔
379	multioutput = True	9✔
380	nout = y.shape[1]	9✔
381	else:
382	raise NotImplementedError("y shape not supported!")	×
383
384	if denoise:	9✔
385	if weights is not None:	9✔
386	raise NotImplementedError(	×
387	"No weights for denoising - the weights are learned."
388	)
389	if Xresampled is not None and selection is not None:	9✔
390	# Select among only the selected features:
391	Xresampled = Xresampled[:, selection]	×
392	if multioutput:	9✔
393	y = np.stack(	9✔
394	[_denoise(X, y[:, i], Xresampled=Xresampled)[1] for i in range(nout)],
395	axis=1,
396	)
397	if Xresampled is not None:	9✔
398	X = Xresampled	×
399	else:
400	X, y = _denoise(X, y, Xresampled=Xresampled)	×
401
402	julia_project = _get_julia_project(julia_project)	9✔
403
404	tmpdir = Path(tempfile.mkdtemp(dir=tempdir))	9✔
405
406	if temp_equation_file:	9✔
407	equation_file = tmpdir / "hall_of_fame.csv"	9✔
408	elif equation_file is None:	9✔
409	date_time = datetime.now().strftime("%Y-%m-%d_%H%M%S.%f")[:-3]	9✔
410	equation_file = "hall_of_fame_" + date_time + ".csv"	9✔
411
412	_create_inline_operators(	9✔
413	binary_operators=binary_operators, unary_operators=unary_operators
414	)
415	_handle_constraints(	9✔
416	binary_operators=binary_operators,
417	unary_operators=unary_operators,
418	constraints=constraints,
419	)
420
421	una_constraints = [constraints[op] for op in unary_operators]	9✔
422	bin_constraints = [constraints[op] for op in binary_operators]	9✔
423
424	try:	9✔
425	term_width = shutil.get_terminal_size().columns	9✔
NEW 426	except:	×
NEW 427	_, term_width = subprocess.check_output(["stty", "size"]).split()	×
428
429	if not already_ran:	9✔
430	from julia import Pkg	9✔
431
432	Pkg.activate(f"{_escape_filename(julia_project)}")	9✔
433	if update:	9✔
434	try:	9✔
435	Pkg.resolve()	9✔
NEW 436	except RuntimeError as e:	×
NEW 437	raise ImportError(	×
438	f"""
439	Required dependencies are not installed or built. Run the following code in the Python REPL:
440
441	>>> import pysr
442	>>> pysr.install()
443
444	Tried to activate project {julia_project} but failed."""
445	) from e
446	Main.eval("using SymbolicRegression")	9✔
447
448	Main.plus = Main.eval("(+)")	9✔
449	Main.sub = Main.eval("(-)")	9✔
450	Main.mult = Main.eval("(*)")	9✔
451	Main.pow = Main.eval("(^)")	9✔
452	Main.div = Main.eval("(/)")	9✔
453
454	Main.custom_loss = Main.eval(loss)	9✔
455
456	mutationWeights = [	9✔
457	float(weightMutateConstant),
458	float(weightMutateOperator),
459	float(weightAddNode),
460	float(weightInsertNode),
461	float(weightDeleteNode),
462	float(weightSimplify),
463	float(weightRandomize),
464	float(weightDoNothing),
465	]
466
467	options = Main.Options(	9✔
468	binary_operators=Main.eval(str(tuple(binary_operators)).replace("'", "")),
469	unary_operators=Main.eval(str(tuple(unary_operators)).replace("'", "")),
470	bin_constraints=bin_constraints,
471	una_constraints=una_constraints,
472	parsimony=float(parsimony),
473	loss=Main.custom_loss,
474	alpha=float(alpha),
475	maxsize=int(maxsize),
476	maxdepth=int(maxdepth),
477	fast_cycle=fast_cycle,
478	migration=migration,
479	hofMigration=hofMigration,
480	fractionReplacedHof=float(fractionReplacedHof),
481	shouldOptimizeConstants=shouldOptimizeConstants,
482	hofFile=_escape_filename(equation_file),
483	npopulations=int(populations),
484	optimizer_algorithm=optimizer_algorithm,
485	optimizer_nrestarts=int(optimizer_nrestarts),
486	optimize_probability=float(optimize_probability),
487	optimizer_iterations=int(optimizer_iterations),
488	perturbationFactor=float(perturbationFactor),
489	annealing=annealing,
490	batching=batching,
491	batchSize=int(min([batchSize, len(X)]) if batching else len(X)),
492	mutationWeights=mutationWeights,
493	warmupMaxsizeBy=float(warmupMaxsizeBy),
494	useFrequency=useFrequency,
495	npop=int(npop),
496	ns=int(tournament_selection_n),
497	probPickFirst=float(tournament_selection_p),
498	ncyclesperiteration=int(ncyclesperiteration),
499	fractionReplaced=float(fractionReplaced),
500	topn=int(topn),
501	verbosity=int(verbosity),
502	progress=progress,
503	terminal_width=int(term_width),
504	**kwargs,
505	)
506
507	np_dtype = {16: np.float16, 32: np.float32, 64: np.float64}[precision]	9✔
508
509	Main.X = np.array(X, dtype=np_dtype).T	9✔
510	if len(y.shape) == 1:	9✔
511	Main.y = np.array(y, dtype=np_dtype)	9✔
512	else:
513	Main.y = np.array(y, dtype=np_dtype).T	9✔
514	if weights is not None:	9✔
515	if len(weights.shape) == 1:	9✔
NEW 516	Main.weights = np.array(weights, dtype=np_dtype)	×
517	else:
518	Main.weights = np.array(weights, dtype=np_dtype).T	9✔
519	else:
520	Main.weights = None	9✔
521
522	cprocs = 0 if multithreading else procs	9✔
523
524	raw_julia_output = Main.EquationSearch(	9✔
525	Main.X,
526	Main.y,
527	weights=Main.weights,
528	niterations=int(niterations),
529	varMap=variable_names,
530	options=options,
531	numprocs=int(cprocs),
532	multithreading=bool(multithreading),
533	)
534
535	_set_globals(	9✔
536	X=X,
537	equation_file=equation_file,
538	variable_names=variable_names,
539	extra_sympy_mappings=extra_sympy_mappings,
540	extra_torch_mappings=extra_torch_mappings,
541	extra_jax_mappings=extra_jax_mappings,
542	output_jax_format=output_jax_format,
543	output_torch_format=output_torch_format,
544	multioutput=multioutput,
545	nout=nout,
546	selection=selection,
547	raw_julia_output=raw_julia_output,
548	)
549
550	equations = get_hof(	9✔
551	equation_file=equation_file,
552	n_features=X.shape[1],
553	variable_names=variable_names,
554	output_jax_format=output_jax_format,
555	output_torch_format=output_torch_format,
556	selection=selection,
557	extra_sympy_mappings=extra_sympy_mappings,
558	extra_jax_mappings=extra_jax_mappings,
559	extra_torch_mappings=extra_torch_mappings,
560	multioutput=multioutput,
561	nout=nout,
562	)
563
564	if delete_tempfiles:	9✔
565	shutil.rmtree(tmpdir)	9✔
566
567	return equations	9✔
568
569
570	def _set_globals(	9✔
571	*,
572	X,
573	equation_file,
574	variable_names,
575	extra_sympy_mappings,
576	extra_torch_mappings,
577	extra_jax_mappings,
578	output_jax_format,
579	output_torch_format,
580	multioutput,
581	nout,
582	selection,
583	raw_julia_output,
584	):
585	global global_state
586
587	global_state["n_features"] = X.shape[1]	9✔
588	global_state["equation_file"] = equation_file	9✔
589	global_state["variable_names"] = variable_names	9✔
590	global_state["extra_sympy_mappings"] = extra_sympy_mappings	9✔
591	global_state["extra_torch_mappings"] = extra_torch_mappings	9✔
592	global_state["extra_jax_mappings"] = extra_jax_mappings	9✔
593	global_state["output_jax_format"] = output_jax_format	9✔
594	global_state["output_torch_format"] = output_torch_format	9✔
595	global_state["multioutput"] = multioutput	9✔
596	global_state["nout"] = nout	9✔
597	global_state["selection"] = selection	9✔
598	global_state["raw_julia_output"] = raw_julia_output	9✔
599
600
601	def _handle_constraints(binary_operators, unary_operators, constraints):	9✔
602	for op in unary_operators:	9✔
603	if op not in constraints:	9✔
604	constraints[op] = -1	9✔
605	for op in binary_operators:	9✔
606	if op not in constraints:	9✔
607	constraints[op] = (-1, -1)	9✔
608	if op in ["plus", "sub"]:	9✔
609	if constraints[op][0] != constraints[op][1]:	9✔
610	raise NotImplementedError(	×
611	"You need equal constraints on both sides for - and *, due to simplification strategies."
612	)
613	elif op == "mult":	9✔
614	# Make sure the complex expression is in the left side.
615	if constraints[op][0] == -1:	×
616	continue	×
617	if constraints[op][1] == -1 or constraints[op][0] < constraints[op][1]:	×
618	constraints[op][0], constraints[op][1] = (	×
619	constraints[op][1],
620	constraints[op][0],
621	)
622
623
624	def _create_inline_operators(binary_operators, unary_operators):	9✔
625	for op_list in [binary_operators, unary_operators]:	9✔
626	for i, op in enumerate(op_list):	9✔
627	is_user_defined_operator = "(" in op	9✔
628
629	if is_user_defined_operator:	9✔
630	Main.eval(op)	9✔
631	# Cut off from the first non-alphanumeric char:
632	first_non_char = [	9✔
633	j
634	for j, char in enumerate(op)
635	if not (char.isalpha() or char.isdigit())
636	][0]
637	function_name = op[:first_non_char]	9✔
638	op_list[i] = function_name	9✔
639
640
641	def _handle_feature_selection(	9✔
642	X, select_k_features, use_custom_variable_names, variable_names, y
643	):
644	if select_k_features is not None:	9✔
645	selection = run_feature_selection(X, y, select_k_features)	9✔
646	print(f"Using features {selection}")	9✔
647	X = X[:, selection]	9✔
648
649	if use_custom_variable_names:	9✔
650	variable_names = [variable_names[i] for i in selection]	9✔
651	else:
652	selection = None	9✔
653	return X, variable_names, selection	9✔
654
655
656	def _check_assertions(	9✔
657	X,
658	binary_operators,
659	unary_operators,
660	use_custom_variable_names,
661	variable_names,
662	weights,
663	y,
664	):
665	# Check for potential errors before they happen
666	assert len(unary_operators) + len(binary_operators) > 0	9✔
667	assert len(X.shape) == 2	9✔
668	assert len(y.shape) in [1, 2]	9✔
669	assert X.shape[0] == y.shape[0]	9✔
670	if weights is not None:	9✔
671	assert weights.shape == y.shape	9✔
672	assert X.shape[0] == weights.shape[0]	9✔
673	if use_custom_variable_names:	9✔
674	assert len(variable_names) == X.shape[1]	9✔
675
676
677	def run_feature_selection(X, y, select_k_features):	9✔
678	"""Use a gradient boosting tree regressor as a proxy for finding
679	the k most important features in X, returning indices for those
680	features as output."""
681
682	from sklearn.ensemble import RandomForestRegressor	9✔
683	from sklearn.feature_selection import SelectFromModel	9✔
684
685	clf = RandomForestRegressor(n_estimators=100, max_depth=3, random_state=0)	9✔
686	clf.fit(X, y)	9✔
687	selector = SelectFromModel(	9✔
688	clf, threshold=-np.inf, max_features=select_k_features, prefit=True
689	)
690	return selector.get_support(indices=True)	9✔
691
692
693	def get_hof(	9✔
694	equation_file=None,
695	n_features=None,
696	variable_names=None,
697	output_jax_format=None,
698	output_torch_format=None,
699	selection=None,
700	extra_sympy_mappings=None,
701	extra_jax_mappings=None,
702	extra_torch_mappings=None,
703	multioutput=None,
704	nout=None,
705	**kwargs,
706	):
707	"""Get the equations from a hall of fame file. If no arguments
708	entered, the ones used previously from a call to PySR will be used."""
709
710	global global_state
711
712	if equation_file is None:	9✔
713	equation_file = global_state["equation_file"]	9✔
714	if n_features is None:	9✔
715	n_features = global_state["n_features"]	9✔
716	if variable_names is None:	9✔
717	variable_names = global_state["variable_names"]	9✔
718	if extra_sympy_mappings is None:	9✔
719	extra_sympy_mappings = global_state["extra_sympy_mappings"]	9✔
720	if extra_jax_mappings is None:	9✔
721	extra_jax_mappings = global_state["extra_jax_mappings"]	9✔
722	if extra_torch_mappings is None:	9✔
723	extra_torch_mappings = global_state["extra_torch_mappings"]	9✔
724	if output_torch_format is None:	9✔
725	output_torch_format = global_state["output_torch_format"]	9✔
726	if output_jax_format is None:	9✔
727	output_jax_format = global_state["output_jax_format"]	9✔
728	if multioutput is None:	9✔
729	multioutput = global_state["multioutput"]	9✔
730	if nout is None:	9✔
731	nout = global_state["nout"]	9✔
732	if selection is None:	9✔
733	selection = global_state["selection"]	9✔
734
735	global_state["selection"] = selection	9✔
736	global_state["equation_file"] = equation_file	9✔
737	global_state["n_features"] = n_features	9✔
738	global_state["variable_names"] = variable_names	9✔
739	global_state["extra_sympy_mappings"] = extra_sympy_mappings	9✔
740	global_state["extra_jax_mappings"] = extra_jax_mappings	9✔
741	global_state["extra_torch_mappings"] = extra_torch_mappings	9✔
742	global_state["output_torch_format"] = output_torch_format	9✔
743	global_state["output_jax_format"] = output_jax_format	9✔
744	global_state["multioutput"] = multioutput	9✔
745	global_state["nout"] = nout	9✔
746	global_state["selection"] = selection	9✔
747
748	try:	9✔
749	if multioutput:	9✔
750	all_outputs = [	9✔
751	pd.read_csv(str(equation_file) + f".out{i}" + ".bkup", sep="\|")
752	for i in range(1, nout + 1)
753	]
754	else:
755	all_outputs = [pd.read_csv(str(equation_file) + ".bkup", sep="\|")]	9✔
756	except FileNotFoundError:	×
757	raise RuntimeError(	×
758	"Couldn't find equation file! The equation search likely exited before a single iteration completed."
759	)
760
761	ret_outputs = []	9✔
762
763	for output in all_outputs:	9✔
764
765	scores = []	9✔
766	lastMSE = None	9✔
767	lastComplexity = 0	9✔
768	sympy_format = []	9✔
769	lambda_format = []	9✔
770	if output_jax_format:	9✔
771	jax_format = []	9✔
772	if output_torch_format:	9✔
773	torch_format = []	9✔
774	use_custom_variable_names = len(variable_names) != 0	9✔
775	local_sympy_mappings = {extra_sympy_mappings, sympy_mappings}	9✔
776
777	if use_custom_variable_names:	9✔
778	sympy_symbols = [sympy.Symbol(variable_names[i]) for i in range(n_features)]	9✔
779	else:
780	sympy_symbols = [sympy.Symbol("x%d" % i) for i in range(n_features)]	9✔
781
782	for _, eqn_row in output.iterrows():	9✔
783	eqn = sympify(eqn_row["Equation"], locals=local_sympy_mappings)	9✔
784	sympy_format.append(eqn)	9✔
785
786	# Numpy:
787	lambda_format.append(CallableEquation(sympy_symbols, eqn, selection))	9✔
788
789	# JAX:
790	if output_jax_format:	9✔
791	from .export_jax import sympy2jax	9✔
792
793	func, params = sympy2jax(	9✔
794	eqn,
795	sympy_symbols,
796	selection=selection,
797	extra_jax_mappings=extra_jax_mappings,
798	)
799	jax_format.append({"callable": func, "parameters": params})	9✔
800
801	# Torch:
802	if output_torch_format:	9✔
803	from .export_torch import sympy2torch	9✔
804
805	module = sympy2torch(	9✔
806	eqn,
807	sympy_symbols,
808	selection=selection,
809	extra_torch_mappings=extra_torch_mappings,
810	)
811	torch_format.append(module)	9✔
812
813	curMSE = eqn_row["MSE"]	9✔
814	curComplexity = eqn_row["Complexity"]	9✔
815
816	if lastMSE is None:	9✔
817	cur_score = 0.0	9✔
818	else:
819	if curMSE > 0.0:	9✔
820	cur_score = -np.log(curMSE / lastMSE) / (	9✔
821	curComplexity - lastComplexity
822	)
823	else:
824	cur_score = np.inf	×
825
826	scores.append(cur_score)	9✔
827	lastMSE = curMSE	9✔
828	lastComplexity = curComplexity	9✔
829
830	output["score"] = np.array(scores)	9✔
831	output["sympy_format"] = sympy_format	9✔
832	output["lambda_format"] = lambda_format	9✔
833	output_cols = [	9✔
834	"Complexity",
835	"MSE",
836	"score",
837	"Equation",
838	"sympy_format",
839	"lambda_format",
840	]
841	if output_jax_format:	9✔
842	output_cols += ["jax_format"]	9✔
843	output["jax_format"] = jax_format	9✔
844	if output_torch_format:	9✔
845	output_cols += ["torch_format"]	9✔
846	output["torch_format"] = torch_format	9✔
847
848	ret_outputs.append(output[output_cols])	9✔
849
850	if multioutput:	9✔
851	return ret_outputs	9✔
852	return ret_outputs[0]	9✔
853
854
855	def best_row(equations=None):	9✔
856	"""Return the best row of a hall of fame file using the score column.
857	By default this uses the last equation file.
858	"""
859	if equations is None:	9✔
860	equations = get_hof()	×
861	if isinstance(equations, list):	9✔
862	return [eq.iloc[np.argmax(eq["score"])] for eq in equations]	9✔
863	return equations.iloc[np.argmax(equations["score"])]	9✔
864
865
866	def best_tex(equations=None):	9✔
867	"""Return the equation with the best score, in latex format
868	By default this uses the last equation file.
869	"""
870	if equations is None:	9✔
871	equations = get_hof()	9✔
872	if isinstance(equations, list):	9✔
873	return [	×
874	sympy.latex(best_row(eq)["sympy_format"].simplify()) for eq in equations
875	]
876	return sympy.latex(best_row(equations)["sympy_format"].simplify())	9✔
877
878
879	def best(equations=None):	9✔
880	"""Return the equation with the best score, in sympy format.
881	By default this uses the last equation file.
882	"""
883	if equations is None:	9✔
884	equations = get_hof()	9✔
885	if isinstance(equations, list):	9✔
886	return [best_row(eq)["sympy_format"].simplify() for eq in equations]	×
887	return best_row(equations)["sympy_format"].simplify()	9✔
888
889
890	def best_callable(equations=None):	9✔
891	"""Return the equation with the best score, in callable format.
892	By default this uses the last equation file.
893	"""
894	if equations is None:	9✔
895	equations = get_hof()	9✔
896	if isinstance(equations, list):	9✔
897	return [best_row(eq)["lambda_format"] for eq in equations]	9✔
898	return best_row(equations)["lambda_format"]	9✔
899
900
901	def _escape_filename(filename):	9✔
902	"""Turns a file into a string representation with correctly escaped backslashes"""
903	str_repr = str(filename)	9✔
904	str_repr = str_repr.replace("\\", "\\\\")	9✔
905	return str_repr	9✔
906
907
908	# https://gist.github.com/garrettdreyfus/8153571
909	def _yesno(question):	9✔
910	"""Simple Yes/No Function."""
911	prompt = f"{question} (y/n): "	9✔
912	ans = input(prompt).strip().lower()	9✔
913	if ans not in ["y", "n"]:	9✔
914	print(f"{ans} is invalid, please try again...")	×
915	return _yesno(question)	×
916	if ans == "y":	9✔
917	return True	9✔
918	return False	9✔
919
920
921	def _denoise(X, y, Xresampled=None):	9✔
922	"""Denoise the dataset using a Gaussian process"""
923	from sklearn.gaussian_process import GaussianProcessRegressor	9✔
924	from sklearn.gaussian_process.kernels import RBF, WhiteKernel, ConstantKernel	9✔
925
926	gp_kernel = RBF(np.ones(X.shape[1])) + WhiteKernel(1e-1) + ConstantKernel()	9✔
927	gpr = GaussianProcessRegressor(kernel=gp_kernel, n_restarts_optimizer=50)	9✔
928	gpr.fit(X, y)	9✔
929	if Xresampled is not None:	9✔
930	return Xresampled, gpr.predict(Xresampled)	×
931
932	return X, gpr.predict(X)	9✔
933
934
935	class CallableEquation:	9✔
936	"""Simple wrapper for numpy lambda functions built with sympy"""
937
938	def __init__(self, sympy_symbols, eqn, selection=None):	9✔
939	self._sympy = eqn	9✔
940	self._sympy_symbols = sympy_symbols	9✔
941	self._selection = selection	9✔
942	self._lambda = lambdify(sympy_symbols, eqn)	9✔
943
944	def __repr__(self):	9✔
945	return f"PySRFunction(X=>{self._sympy})"	9✔
946
947	def __call__(self, X):	9✔
948	if self._selection is not None:	9✔
949	return self._lambda(*X[:, self._selection].T)	×
950	return self._lambda(*X.T)	9✔
951
952
953	def _get_julia_project(julia_project):	9✔
954	pkg_directory = Path(__file__).parents[1]	9✔
955	if julia_project is None:	9✔
956	return pkg_directory	9✔
NEW 957	return Path(julia_project)	×
958
959
960	def silence_julia_warning():	9✔
961	global is_julia_warning_silenced
NEW 962	is_julia_warning_silenced = True	×
963
964
965	def init_julia():	9✔
966	"""Initialize julia binary, turning off compiled modules if needed."""
967	global is_julia_warning_silenced
968	from julia.core import JuliaInfo, UnsupportedPythonError	9✔
969
970	info = JuliaInfo.load(julia="julia")	9✔
971	if not info.is_pycall_built():	9✔
NEW 972	raise ImportError(	×
973	"""
974	Required dependencies are not installed or built. Run the following code in the Python REPL:
975
976	>>> import pysr
977	>>> pysr.install()"""
978	)
979
980	Main = None	9✔
981	try:	9✔
982	from julia import Main as _Main	9✔
983
984	Main = _Main	9✔
NEW 985	except UnsupportedPythonError:	×
NEW 986	if not is_julia_warning_silenced:	×
NEW 987	warnings.warn(	×
988	"""
989	Your Python version is statically linked to libpython. For example, this could be the python included with conda, or maybe your system's built-in python.
990	This will still work, but the precompilation cache for Julia will be turned off, which may result in slower startup times on the initial pysr() call.
991
992	To install a Python version that is dynamically linked to libpython, pyenv is recommended (https://github.com/pyenv/pyenv).
993
994	To silence this warning, you can run pysr.silence_julia_warning() after importing pysr."""
995	)
NEW 996	from julia.core import Julia	×
997
NEW 998	jl = Julia(compiled_modules=False)	×
NEW 999	from julia import Main as _Main	×
1000
NEW 1001	Main = _Main	×
1002
1003	return Main	9✔

MilesCranmer / PySR / 1725881820

Source File Press 'n' to go to next uncovered line, 'b' for previous

Source File
Press 'n' to go to next uncovered line, 'b' for previous