Installation instructions

[lollcat]

Hi!

I'm having some issues with the install - specifically I'm getting an error from the use of jax_dataclasses within flox (described lower down in this issue). Would it be possible to please provide more detailed setup instructions? I think this error would be resolved if I use the same versions of all libraries as what was used for the original code.

Below I summarise how I am currently setting up the repo, and the error I get:

Install

1. Make venv with python 3.11
2. Clone repo

git clone https://github.com/noegroup/rigid-flows.git
cd rigid-flows

3. Install latest jax (cpu)

pip install --upgrade pip
pip install --upgrade "jax[cpu]"

4. Install custom libraries not on pypi

git clone https://github.com/jonkhler/mlparams.git
pip install -e ./mlparams/
git clone https://github.com/noegroup/flox.git
pip install -e ./flox/

Side comment: I think it should be noted in the README that these libraries need to be installed like this. Currently these libraries are in the setup.py file so if a user naively tries to setup the repo by cloning it and running pip install -e . then this will not work.

5. Install other required libraries

pip install genice2 pymbar GitPython tensorflow mdtraj openmm

The error

Then if I run python experiments/run.py I get the following error:

  File "/home/laurence/Documents/work/code/rigid-flows/experiments/run.py", line 13, in <module>
    from flox._src.flow.api import Transform
  File "/home/laurence/Documents/work/code/rigid-flows/flox/flox/__init__.py", line 21, in <module>
    from . import bulk, flow, geom, nn, util
  File "/home/laurence/Documents/work/code/rigid-flows/flox/flox/flow.py", line 1, in <module>
    from ._src.flow import *
  File "/home/laurence/Documents/work/code/rigid-flows/flox/flox/_src/flow/__init__.py", line 3, in <module>
    from .sampling import *
  File "/home/laurence/Documents/work/code/rigid-flows/flox/flox/_src/flow/sampling.py", line 30, in <module>
    @pytree_dataclass(frozen=True)
     ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/home/laurence/miniconda3/envs/rigid_flows/lib/python3.11/site-packages/jax_dataclasses/_dataclasses.py", line 38, in wrap
    return _register_pytree_dataclass(dataclasses.dataclass(cls, **kwargs))
                                      ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/home/laurence/miniconda3/envs/rigid_flows/lib/python3.11/dataclasses.py", line 1230, in dataclass
    return wrap(cls)
           ^^^^^^^^^
  File "/home/laurence/miniconda3/envs/rigid_flows/lib/python3.11/dataclasses.py", line 1220, in wrap
    return _process_class(cls, init, repr, eq, order, unsafe_hash,
           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/home/laurence/miniconda3/envs/rigid_flows/lib/python3.11/dataclasses.py", line 958, in _process_class
    cls_fields.append(_get_field(cls, name, type, kw_only))
                      ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/home/laurence/miniconda3/envs/rigid_flows/lib/python3.11/dataclasses.py", line 815, in _get_field
    raise ValueError(f'mutable default {type(f.default)} for field '
ValueError: mutable default <class 'jaxlib.xla_extension.ArrayImpl'> for field ldj is not allowed: use default_factory