milescranmer / symbolic_deep_learning

Code for "Discovering Symbolic Models from Deep Learning with Inductive Biases"

Discovering Symbolic Models from Deep Learning with Inductive Biases

This repository is the official implementation of Discovering Symbolic Models from Deep Learning with Inductive Biases.

Miles Cranmer, Alvaro Sanchez-Gonzalez, Peter Battaglia, Rui Xu, Kyle Cranmer, David Spergel, Shirley Ho

Check out our Blog, Paper, Video, and Interactive Demo.

Requirements

For model:

• pytorch
• pytorch-geometric
• numpy

Symbolic regression:

• PySR, our new open-source Eureqa alternative

For simulations:

• jax (simple N-body simulations)
• quijote (Dark matter data; optional)
• tqdm
• matplotlib

Training

To train an example model from the paper, try out the demo.

Full model definitions are given in models.py. Data is generated from simulate.py.

Results

We train on simulations produced by the following equations: giving us time series:

We recorded performance for each model: and also measured how well each model's messages correlated with a linear combination of forces:

Finally, we trained on a dark matter simulation and extracted the following equations from the message function:

symbolic_deep_learning's Issues

Cleanup

Dear @MilesCranmer,

Awesome work with this repo.

I really like how you are using jax to generate the data and PyTorch Geometric for the model.
I think you could have made your code even cleaner if you used PyTorch Lightning.
We worked with Matthias Fey to better support PyTorch Geometric and working DGL team right now for the same purpose.

Find examples there: https://github.com/rusty1s/pytorch_geometric/blob/master/examples/pytorch_lightning/graph_sage.py#L72

Best,
T.C

Demo not working due to bad torch version

Hi,

I tried to run your demo and failed when importing models with,

> ---------------------------------------------------------------------------
> 
> OSError                                   Traceback (most recent call last)
> 
> <ipython-input-5-bd1aa79e6e09> in <module>()
> ----> 1 import models
>       2 import simulate
> 
> /content/models.py in <module>()
>       4 from torch.functional import F
>       5 from torch.optim import Adam
> ----> 6 from torch_geometric.nn import MetaLayer, MessagePassing
>       7 from torch.nn import Sequential as Seq, Linear as Lin, ReLU, Softplus
>       8 from torch.autograd import Variable, grad
> 
> /usr/local/lib/python3.6/dist-packages/torch_geometric/__init__.py in <module>()
>       1 from .debug import is_debug_enabled, debug, set_debug
> ----> 2 import torch_geometric.nn
>       3 import torch_geometric.data
>       4 import torch_geometric.datasets
>       5 import torch_geometric.transforms
> 
> /usr/local/lib/python3.6/dist-packages/torch_geometric/nn/__init__.py in <module>()
>       1 from .meta import MetaLayer
> ----> 2 from .data_parallel import DataParallel
>       3 from .reshape import Reshape
>       4 from .conv import *  # noqa
>       5 from .norm import *  # noqa
> 
> /usr/local/lib/python3.6/dist-packages/torch_geometric/nn/data_parallel.py in <module>()
>       3 
>       4 import torch
> ----> 5 from torch_geometric.data import Batch
>       6 
>       7 
> 
> /usr/local/lib/python3.6/dist-packages/torch_geometric/data/__init__.py in <module>()
> ----> 1 from .data import Data
>       2 from .batch import Batch
>       3 from .dataset import Dataset
>       4 from .in_memory_dataset import InMemoryDataset
>       5 from .dataloader import DataLoader, DataListLoader, DenseDataLoader
> 
> /usr/local/lib/python3.6/dist-packages/torch_geometric/data/data.py in <module>()
>       6 import torch
>       7 import torch_geometric
> ----> 8 from torch_sparse import coalesce, SparseTensor
>       9 from torch_geometric.utils import (contains_isolated_nodes,
>      10                                    contains_self_loops, is_undirected)
> 
> /usr/local/lib/python3.6/dist-packages/torch_sparse/__init__.py in <module>()
>      11 ]:
>      12     torch.ops.load_library(importlib.machinery.PathFinder().find_spec(
> ---> 13         library, [osp.dirname(__file__)]).origin)
>      14 
>      15 if torch.cuda.is_available() and torch.version.cuda:  # pragma: no cover
> 
> /usr/local/lib/python3.6/dist-packages/torch/_ops.py in load_library(self, path)
>     103             # static (global) initialization code in order to register custom
>     104             # operators with the JIT.
> --> 105             ctypes.CDLL(path)
>     106         self.loaded_libraries.add(path)
>     107 
> 
> /usr/lib/python3.6/ctypes/__init__.py in __init__(self, name, mode, handle, use_errno, use_last_error)
>     346 
>     347         if handle is None:
> --> 348             self._handle = _dlopen(self._name, mode)
>     349         else:
>     350             self._handle = handle
> 
> OSError: /usr/local/lib/python3.6/dist-packages/torch_sparse/_convert.so: undefined symbol: _ZN3c104impl23ExcludeDispatchKeyGuardC1ENS_11DispatchKeyE

I've solved it changing torch version,

!export CUDA=cu101 && pip install --upgrade torch-scatter==latest+${CUDA} torch-sparse==latest+${CUDA} -f https://pytorch-geometric.com/whl/torch-1.6.0.html

to 1.7.0,

!export CUDA=cu101 && pip install --upgrade torch-scatter==latest+${CUDA} torch-sparse==latest+${CUDA} -f https://pytorch-geometric.com/whl/torch-1.7.0.html

Regards,

Problems about GPU utilisation rate

Hi Miles,

I am running ur codes on GPU but I found the utilization rate is usually very low like around 30% or below. Did u meet similar questions before and have some suggestions? Is this because that we use Jax package? Kindly let me know.

Many thanks!

Best,
Haiwei

PyTorch and torch_sparse versions do not match.

Hello again Miles,

Running your demo as it is has the following issues,

First it takes a lot of time in building wheels for torch-sparse, torch-scatter and torch geometric. I don't know if this is normal, but at the end it finishes, most of the times.

The problem is then when the code imports models.py, this is the error:

---------------------------------------------------------------------------

RuntimeError                              Traceback (most recent call last)

<ipython-input-11-406efe1c43e4> in <module>()
      1 import simulate
----> 2 import models

4 frames

/content/models.py in <module>()
      4 from torch.functional import F
      5 from torch.optim import Adam
----> 6 from torch_geometric.nn import MetaLayer, MessagePassing
      7 from torch.nn import Sequential as Seq, Linear as Lin, ReLU, Softplus
      8 from torch.autograd import Variable, grad

/usr/local/lib/python3.7/dist-packages/torch_geometric/__init__.py in <module>()
      3 
      4 from .debug import is_debug_enabled, debug, set_debug
----> 5 import torch_geometric.data
      6 import torch_geometric.transforms
      7 import torch_geometric.utils

/usr/local/lib/python3.7/dist-packages/torch_geometric/data/__init__.py in <module>()
----> 1 from .data import Data
      2 from .temporal import TemporalData
      3 from .batch import Batch
      4 from .dataset import Dataset
      5 from .in_memory_dataset import InMemoryDataset

/usr/local/lib/python3.7/dist-packages/torch_geometric/data/data.py in <module>()
      6 import torch
      7 import torch_geometric
----> 8 from torch_sparse import coalesce, SparseTensor
      9 from torch_geometric.utils import (contains_isolated_nodes,
     10                                    contains_self_loops, is_undirected)

/usr/local/lib/python3.7/dist-packages/torch_sparse/__init__.py in <module>()
     28     if t_major != major:
     29         raise RuntimeError(
---> 30             f'Detected that PyTorch and torch_sparse were compiled with '
     31             f'different CUDA versions. PyTorch has CUDA version '
     32             f'{t_major}.{t_minor} and torch_sparse has CUDA version '

RuntimeError: Detected that PyTorch and torch_sparse were compiled with different CUDA versions. PyTorch has CUDA version 10.1 and torch_sparse has CUDA version 11.0. Please reinstall the torch_sparse that matches your PyTorch install.

I've tried playing with the versions to see if I could fix it but no luck.

Thanks,

Training Data Generation for Dark Matter Cosmology Example

Dear @MilesCranmer and team,
Thank you for providing this great framework in an open-source way. Since your framework showed great performance on the dark matter cosmology example, I wonder about some implementation details of the same, which were not clear to me after reading the paper and going through the spring-mass example available in this repository.

The full halo graph has a huge number of nodes and edges, and I am unsure how this is reflected in the model. I am aware that each halo considers neighbors within a certain radius. But how do you generate training data from this huge graph? Do you extract the features of one halo and its neighbors at a time to build the training data? Also, as far as I understand, the cosmology example does not have a time dependence as in the spring example. How is this reflected in the model?

I'm really looking forward to finding out more about your implementation details.

Thank you in advance for your time!

Support for Windows systems that cannot install jax?

Hi, thank you for your amazing work on this and PySR and for making it available for us to use.

I have encountered a problem with the demo file because it requires jaxlib, which isn't available on windows. Additionally, I don't have a gpu either. Based on my understanding, jax is only being used to generate the data. However I proficient enough to be able to understand the code and write a general version that works with other types of data so that jaxlib isn't required, without first being able to run the entire demo to understand each step in detail

Would you please advice me on what I can do to make this work on windows?
Thank you very much.
Regards,
jstan

Getting back 0 std messages from own data

Dear @MilesCranmer and team,

I am using this repo for a school project and am experiencing an annoying issue. When I implement my own data, I sometimes get back 0 std messages. All e(number) messages have the same value per column. I've just ran the same code twice in a row and once I got back all 0's and now I do get back non-zero numbers.

These are from small tests with only 1 epoch, but when I ran the code for 50 epochs earlier I got back 0's. This is very annoying since I always have to wait a long time for the code to finish.

Is there something I could do to avoid this issue? Maybe something to do with the data itself? Or a hyperparameter?

Best regards,
L.W.