Compatibility with newer TF (2.11-2.15) (#271)

* Fixed import bugs with newer TensorFlow versions.
* Minor refactor to help with future upgrade to TF 2.16.
* Added message for TF >2.15 not being supported.
* Updated readme.
* Updated installation instructions to include TF versions.

README.md

@@ -0,0 +1,120 @@
+# osl-dynamics
+
+See the read the docs page for a description of this project: [https://osl-dynamics.readthedocs.io](https://osl-dynamics.readthedocs.io).
+
+## Citation
+
+If you find this toolbox useful, please cite:
+
+> **Chetan Gohil, Rukuang Huang, Evan Roberts, Mats WJ van Es, Andrew J Quinn, Diego Vidaurre, Mark W Woolrich (2024) osl-dynamics, a toolbox for modeling fast dynamic brain activity eLife 12:RP91949.**
+
+## Installation
+
+### Conda
+
+We recommend installing osl-dynamics within a virtual environment. You can do this with [Anaconda](https://docs.anaconda.com/free/anaconda/install/index.html) (or [miniconda](https://docs.conda.io/projects/miniconda/en/latest/miniconda-install.html).
+
+Below we describe how to install osl-dynamics from source. We recommend using the conda environment files in `/envs`.
+
+### Linux
+
+```
+git clone https://github.com/OHBA-analysis/osl-dynamics.git
+cd osl-dynamics
+conda env create -f envs/linux.yml
+conda activate osld
+pip install -e .
+```
+
+### Mac
+
+For a Mac, the installation of TensorFlow is slightly different to a Linux computer. We recommend using the lines above replacing the Linux environment file `envs/linux.yml` with the Mac environment file `envs/mac.yml`.
+
+### Windows
+
+If you are using a Windows computer, we recommend first installing Linux (Ubuntu) as a Windows Subsystem by following the instructions [here](https://ubuntu.com/wsl). Then following the instructions above in the Ubuntu terminal.
+
+### Within an osl environment
+
+If you have already installed [OSL](https://github.com/OHBA-analysis/osl) you can install osl-dynamics in the `osl` environment with:
+
+```
+conda activate osl
+cd osl-dynamics
+pip install tensorflow==2.11.0
+pip install tensorflow-probability==0.19.0
+pip install -e .
+```
+
+Note, if you're using a Mac computer you need to install TensorFlow with `pip install tensorflow-macos==2.11.0` instead of `tensorflow==2.11.0`.
+
+### TensorFlow versions
+
+osl-dynamics can be used with the following TensorFlow versions:
+
+| tensorflow  | tensorflow-probability |
+| ------------- | ------------- |
+| 2.11 | 0.19  |
+| 2.12 | 0.19  |
+| 2.13 | 0.20  |
+| 2.14 | 0.21  |
+| 2.15 | 0.22  |
+
+### Removing osl-dynamics
+
+Simply delete the conda environment and repository:
+
+```
+conda env remove -n osld
+rm -rf osl-dynamics
+```
+
+## Documentation
+
+The read the docs page should be automatically updated whenever there's a new commit on the `main` branch.
+
+The documentation is included as docstrings in the source code. Please write docstrings to any classes or functions you add following the [numpy style](https://numpydoc.readthedocs.io/en/latest/format.html). The API reference documentation will only be automatically generated if the docstrings are written correctly. The documentation directory `/doc` also contains `.rst` files that provide additional info regarding installation, development, the models, etc.
+
+To compile the documentation locally you need to install the required packages (sphinx, etc.) in your conda environment:
+
+```
+cd osl-dynamics
+pip install -r doc/requirements.txt
+```
+
+To compile the documentation locally use:
+
+```
+python setup.py build_sphinx
+```
+
+The local build of the documentation webpage can be found in `build/sphinx/html/index.html`.
+
+## Releases
+
+A couple packages are needed to build and upload a project to PyPI, these can be installed in your conda environment with:
+
+```
+pip install build twine
+```
+
+The following steps can be used to release a new version:
+
+1. Update the version on line 5 of `setup.cfg` by removing `dev` from the version number.
+2. Commit the updated setup.cfg to the `main` branch of the GitHub repo.
+3. Delete any old distributions that have been built (if there are any): `rm -r dist`.
+4. Build a distribution in the osl-dynamics root directory with `python -m build`. This will create a new directory called `dist`.
+5. Test the build by installing in a test conda environment with `cd dist; pip install <build>.whl`.
+6. Upload the distribution to PyPI with `twine upload dist/*`. You will need to enter the username and password that you used to register with [https://pypi.org](https://pypi.org).
+7. Tag the commit uploaded to PyPI with the version number using the 'Create a new release' link on the right of the GitHub repo webpage.
+8. Change the version to `X.Y.devZ` in `setup.cfg` and commit the new dev version to `main`.
+
+The uploaded distribution will then be available to be installed with:
+
+```
+pip install osl-dynamics
+```
+
+## Editing Source Code
+
+See [here](https://github.com/OHBA-analysis/osl-dynamics/blob/main/doc/using_bmrc.rst) for useful info regarding how to use the Oxford BMRC cluster and how to edit the source code.

doc/install.rst

@@ -63,6 +63,28 @@ Windows Instructions
 
 If you are using a Windows computer, we recommend first installing linux (Ubuntu) as a Windows Subsystem by following the instructions `here <https://ubuntu.com/wsl>`_. Then following the instructions above in the Ubuntu terminal.
 
+TensorFlow Versions
+-------------------
+
+osl-dynamics has been tested with the following versions:
+
+.. list-table::
+   :widths: 25 25
+   :header-rows: 1
+
+   * - tensorflow
+     - tensorflow-probability
+   * - 2.11
+     - 0.19
+   * - 2.12
+     - 0.19
+   * - 2.13
+     - 0.20
+   * - 2.14
+     - 0.22
+   * - 2.15
+     - 0.22
+
 Training Speed
 --------------
 

osl_dynamics/inference/layers.py

@@ -146,6 +146,33 @@ def call(self, inputs, **kwargs):
         return tf.concat(inputs, axis=self.axis)
 
 
+class SplitLayer(layers.Layer):
+    """Wrapper for `tf.split \
+    <https://www.tensorflow.org/api_docs/python/tf/split>`_.
+
+    Parameters
+    ----------
+    num_or_size_splits : int or list
+        Split to apply.
+    axis : int
+        Axis to split along.
+    kwargs : keyword arguments, optional
+        Keyword arguments to pass to the base class.
+    """
+
+    def __init__(self, num_or_size_splits, axis, **kwargs):
+        super().__init__(**kwargs)
+        self.num_or_size_splits = num_or_size_splits
+        self.axis = axis
+
+    def call(self, inputs, **kwargs):
+        return tf.split(
+            inputs,
+            num_or_size_splits=self.num_or_size_splits,
+            axis=self.axis,
+        )
+
+
 class MatMulLayer(layers.Layer):
     """Wrapper for `tf.matmul \
     <https://www.tensorflow.org/api_docs/python/tf/linalg/matmul>`_.

osl_dynamics/inference/optimizers.py

@@ -2,11 +2,19 @@
 
 """
 
+from packaging import version
+
 import tensorflow as tf
-from keras.optimizers.optimizer_v2 import optimizer_v2
+
+if version.parse(tf.__version__) < version.parse("2.12"):
+    from keras.optimizers.optimizer_v2.optimizer_v2 import OptimizerV2 as Optimizer
+elif version.parse(tf.__version__) < version.parse("2.13"):
+    from keras.optimizers.legacy.optimizer_v2 import OptimizerV2 as Optimizer
+else:
+    from keras.optimizers.legacy import Optimizer
 
 
-class ExponentialMovingAverage(optimizer_v2.OptimizerV2):
+class ExponentialMovingAverage(Optimizer):
     """Optimizer for applying a exponential moving average update.
 
     Parameters
@@ -27,7 +35,7 @@ def _resource_apply_dense(self, grad, var):
         return var.assign((1.0 - self.decay) * var + self.decay * grad)
 
 
-class MarkovStateModelOptimizer(optimizer_v2.OptimizerV2):
+class MarkovStateModelOptimizer(Optimizer):
     """Optimizer for a model containing a hidden state Markov chain.
 
     Parameters

osl_dynamics/models/hmm.py

@@ -1660,22 +1660,16 @@ def _model_structure(self):
 
         config = self.config
 
-        # Inputs
+        # Definition of layers
         inputs = layers.Input(
             shape=(config.sequence_length, config.n_channels + config.n_states),
             name="inputs",
         )
-        data, gamma = tf.split(inputs, [config.n_channels, config.n_states], axis=2)
-
-        # Static loss scaling factor
         static_loss_scaling_factor_layer = StaticLossScalingFactorLayer(
             config.sequence_length,
             config.loss_calc,
             name="static_loss_scaling_factor",
         )
-        static_loss_scaling_factor = static_loss_scaling_factor_layer(data)
-
-        # Definition of layers
         means_layer = VectorsLayer(
             config.n_states,
             config.n_channels,
@@ -1712,6 +1706,8 @@ def _model_structure(self):
         )
 
         # Data flow
+        data, gamma = tf.split(inputs, [config.n_channels, config.n_states], axis=2)
+        static_loss_scaling_factor = static_loss_scaling_factor_layer(data)
         mu = means_layer(
             data, static_loss_scaling_factor=static_loss_scaling_factor
         )  # data not used

osl_dynamics/models/hmm_poi.py

@@ -490,14 +490,11 @@ def _model_structure(self):
 
         config = self.config
 
-        # Inputs
+        # Definition of layers
         inputs = layers.Input(
             shape=(config.sequence_length, config.n_channels + config.n_states),
             name="inputs",
         )
-        data, gamma = tf.split(inputs, [config.n_channels, config.n_states], axis=2)
-
-        # Definition of layers
         log_rates_layer = VectorsLayer(
             config.n_states,
             config.n_channels,
@@ -510,6 +507,7 @@ def _model_structure(self):
         )
 
         # Data flow
+        data, gamma = tf.split(inputs, [config.n_channels, config.n_states], axis=2)
         log_rates = log_rates_layer(data)  # data not used
         ll_loss = ll_loss_layer([data, log_rates, gamma, None])