MilesCranmer commited on
Commit
df4b0b3
·
1 Parent(s): 456e76d

Update to SymbolicRegression.jl 0.4.0

Browse files
Files changed (1) hide show
  1. pysr/sr.py +30 -26
pysr/sr.py CHANGED
@@ -103,7 +103,7 @@ def pysr(X=None, y=None, weights=None,
103
  limitPowComplexity=False, #deprecated
104
  threads=None, #deprecated
105
  julia_optimization=3,
106
- local_install=None,
107
  ):
108
  """Run symbolic regression to fit f(X[i, :]) ~ y[i] for all i.
109
  Note: most default parameters have been tuned over several example
@@ -189,9 +189,8 @@ def pysr(X=None, y=None, weights=None,
189
  :param julia_optimization: int, Optimization level (0, 1, 2, 3)
190
  :param tempdir: str or None, directory for the temporary files
191
  :param delete_tempfiles: bool, whether to delete the temporary files after finishing
192
- :param local_install: str or None, the root directory for the
193
- `SymbolicRegression.jl` repo, so one can customize the search code
194
- instead of using the pre-compiled package.
195
  :returns: pd.DataFrame, Results dataframe, giving complexity, MSE, and equations
196
  (as strings).
197
 
@@ -262,7 +261,7 @@ def pysr(X=None, y=None, weights=None,
262
  weightSimplify=weightSimplify,
263
  constraints=constraints,
264
  extra_sympy_mappings=extra_sympy_mappings,
265
- local_install=local_install)
266
 
267
  kwargs = {**_set_paths(tempdir), **kwargs}
268
 
@@ -319,33 +318,29 @@ def _final_pysr_process(julia_optimization, runfile_filename, timeout, **kwargs)
319
 
320
  def _create_julia_files(dataset_filename, def_datasets, hyperparam_filename, def_hyperparams,
321
  fractionReplaced, ncyclesperiteration, niterations, npop,
322
- runfile_filename, topn, verbosity, local_install, procs, weights,
323
  X, variable_names, **kwargs):
324
  with open(hyperparam_filename, 'w') as f:
325
  print(def_hyperparams, file=f)
326
  with open(dataset_filename, 'w') as f:
327
  print(def_datasets, file=f)
328
  with open(runfile_filename, 'w') as f:
329
- print(f'using Distributed', file=f)
330
- print(f'procs = addprocs({procs})', file=f)
331
- if local_install is None:
332
- print(f'@everywhere using SymbolicRegression', file=f)
333
- else:
334
- local_install = Path(local_install) / "src" / "SymbolicRegression.jl"
335
- print(f'@everywhere include("{_escape_filename(local_install)}")', file=f)
336
- print(f'@everywhere using .SymbolicRegression', file=f)
337
- print(f'@everywhere include("{_escape_filename(hyperparam_filename)}")', file=f)
338
- print(f'@everywhere include("{_escape_filename(dataset_filename)}")', file=f)
339
  if len(variable_names) == 0:
340
  varMap = "[" + ",".join([f'"x{i}"' for i in range(X.shape[1])]) + "]"
341
  else:
342
  varMap = "[" + ",".join(variable_names) + "]"
343
 
344
  if weights is not None:
345
- print(f'EquationSearch(X, y, weights=weights, niterations={niterations:d}, varMap={varMap}, options=options)', file=f)
346
  else:
347
- print(f'EquationSearch(X, y, niterations={niterations:d}, varMap={varMap}, options=options)', file=f)
348
- print(f'rmprocs(procs)', file=f)
349
 
350
 
351
  def _make_datasets_julia_str(X, X_filename, weights, weights_filename, y, y_filename, **kwargs):
@@ -371,13 +366,22 @@ def _make_hyperparams_julia_str(X, alpha, annealing, batchSize, batching, binary
371
  ncyclesperiteration, fractionReplaced, topn, verbosity,
372
  weightDeleteNode, weightDoNothing, weightInsertNode, weightMutateConstant,
373
  weightMutateOperator, weightRandomize, weightSimplify, weights, **kwargs):
374
- def_hyperparams += f"""plus=SymbolicRegression.plus
375
- sub=SymbolicRegression.sub
376
- mult=SymbolicRegression.mult
 
 
 
 
 
 
 
 
 
377
  square=SymbolicRegression.square
378
  cube=SymbolicRegression.cube
379
- pow=SymbolicRegression.pow
380
- div=SymbolicRegression.div
381
  logm=SymbolicRegression.logm
382
  logm2=SymbolicRegression.logm2
383
  logm10=SymbolicRegression.logm10
@@ -388,8 +392,8 @@ relu=SymbolicRegression.relu
388
  logical_or=SymbolicRegression.logical_or
389
  logical_and=SymbolicRegression.logical_and
390
 
391
- options = SymbolicRegression.Options(binary_operators={'(' + ', '.join(binary_operators) + ')'},
392
- unary_operators={'(' + ', '.join(unary_operators) + ')'},
393
  {constraints_str}
394
  parsimony={parsimony:f}f0,
395
  alpha={alpha:f}f0,
 
103
  limitPowComplexity=False, #deprecated
104
  threads=None, #deprecated
105
  julia_optimization=3,
106
+ julia_project=None,
107
  ):
108
  """Run symbolic regression to fit f(X[i, :]) ~ y[i] for all i.
109
  Note: most default parameters have been tuned over several example
 
189
  :param julia_optimization: int, Optimization level (0, 1, 2, 3)
190
  :param tempdir: str or None, directory for the temporary files
191
  :param delete_tempfiles: bool, whether to delete the temporary files after finishing
192
+ :param julia_project: str or None, a Julia environment location containing
193
+ a Project.toml (and potentially the source code for SymbolicRegression.jl)
 
194
  :returns: pd.DataFrame, Results dataframe, giving complexity, MSE, and equations
195
  (as strings).
196
 
 
261
  weightSimplify=weightSimplify,
262
  constraints=constraints,
263
  extra_sympy_mappings=extra_sympy_mappings,
264
+ julia_project=julia_project)
265
 
266
  kwargs = {**_set_paths(tempdir), **kwargs}
267
 
 
318
 
319
  def _create_julia_files(dataset_filename, def_datasets, hyperparam_filename, def_hyperparams,
320
  fractionReplaced, ncyclesperiteration, niterations, npop,
321
+ runfile_filename, topn, verbosity, julia_project, procs, weights,
322
  X, variable_names, **kwargs):
323
  with open(hyperparam_filename, 'w') as f:
324
  print(def_hyperparams, file=f)
325
  with open(dataset_filename, 'w') as f:
326
  print(def_datasets, file=f)
327
  with open(runfile_filename, 'w') as f:
328
+ if julia_project is not None:
329
+ julia_project = Path(julia_project)
330
+ print(f'import Pkg', file=f)
331
+ print(f'Pkg.activate("{_escape_filename(julia_project)}")', file=f)
332
+ print(f'using SymbolicRegression', file=f)
333
+ print(f'include("{_escape_filename(hyperparam_filename)}")', file=f)
334
+ print(f'include("{_escape_filename(dataset_filename)}")', file=f)
 
 
 
335
  if len(variable_names) == 0:
336
  varMap = "[" + ",".join([f'"x{i}"' for i in range(X.shape[1])]) + "]"
337
  else:
338
  varMap = "[" + ",".join(variable_names) + "]"
339
 
340
  if weights is not None:
341
+ print(f'EquationSearch(X, y, weights=weights, niterations={niterations:d}, varMap={varMap}, options=options, numprocs={procs})', file=f)
342
  else:
343
+ print(f'EquationSearch(X, y, niterations={niterations:d}, varMap={varMap}, options=options, numprocs={procs})', file=f)
 
344
 
345
 
346
  def _make_datasets_julia_str(X, X_filename, weights, weights_filename, y, y_filename, **kwargs):
 
366
  ncyclesperiteration, fractionReplaced, topn, verbosity,
367
  weightDeleteNode, weightDoNothing, weightInsertNode, weightMutateConstant,
368
  weightMutateOperator, weightRandomize, weightSimplify, weights, **kwargs):
369
+ def tuple_fix(ops):
370
+ if len(ops) > 1:
371
+ return ', '.join(ops)
372
+ elif len(ops) == 0:
373
+ return ''
374
+ else:
375
+ return ops[0] + ','
376
+
377
+ def_hyperparams += f"""\n
378
+ plus=(+)
379
+ sub=(-)
380
+ mult=(*)
381
  square=SymbolicRegression.square
382
  cube=SymbolicRegression.cube
383
+ pow=(^)
384
+ div=(/)
385
  logm=SymbolicRegression.logm
386
  logm2=SymbolicRegression.logm2
387
  logm10=SymbolicRegression.logm10
 
392
  logical_or=SymbolicRegression.logical_or
393
  logical_and=SymbolicRegression.logical_and
394
 
395
+ options = SymbolicRegression.Options(binary_operators={'(' + tuple_fix(binary_operators) + ')'},
396
+ unary_operators={'(' + tuple_fix(unary_operators) + ')'},
397
  {constraints_str}
398
  parsimony={parsimony:f}f0,
399
  alpha={alpha:f}f0,