From 389cfc78f671d5799d2efcdee95baa564975bbbe Mon Sep 17 00:00:00 2001 From: Sampreet Kalita <9553215+Sampreet@users.noreply.github.com> Date: Sat, 10 May 2025 10:02:02 +0100 Subject: [PATCH 1/5] Update Requirements --- .github/workflows/python-tox.yml | 1 - CHANGELOG.md | 6 ++++ CONTRIBUTING.md | 4 +-- README.md | 19 +++++------- pylintrc | 5 ++- pyproject.toml | 5 +-- quantrl/__init__.py | 2 +- quantrl/backends/jax.py | 18 +++++------ quantrl/envs/base.py | 4 +-- quantrl/io.py | 4 +-- quantrl/solvers/measure.py | 53 ++++++++++++++++++-------------- requirements.txt | 3 +- requirements_tox.txt | 9 +++--- 13 files changed, 72 insertions(+), 61 deletions(-) diff --git a/.github/workflows/python-tox.yml b/.github/workflows/python-tox.yml index 333cd2b..bdda448 100644 --- a/.github/workflows/python-tox.yml +++ b/.github/workflows/python-tox.yml @@ -35,7 +35,6 @@ jobs: strategy: matrix: python: [ - "3.10", "3.12", ] diff --git a/CHANGELOG.md b/CHANGELOG.md index 5d3c705..b29e6ec 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -1,5 +1,11 @@ # Changelog +## 2025/05/10 - 00 - v0.0.9 - Update Requirements +* Added support for NumPy version `2.0+` with minimum Python version `3.12`. +* Minor fixes to `quantrl.backends.jax` and `quantrl.io` modules. +* Minor changes to `CONTRIBUTING.md` and `pylintrc`. +* Updated `README`, `pyproject.toml` and `requirements`. + ## 2024/10/14 - 00 - v0.0.8 - Instantiation and GitHub CI * Instantiated backends and solvers with different numerical libraries: * Added `context_manager` modules to `quantrl.backends` and `quantrl.solvers` packages. diff --git a/CONTRIBUTING.md b/CONTRIBUTING.md index b2aa736..72b6e6d 100644 --- a/CONTRIBUTING.md +++ b/CONTRIBUTING.md @@ -57,11 +57,11 @@ ROOT_DIR/ ├───CODE_OF_CONDUCT.md ├───CONTRIBUTING.md ├───LICENSE -├───MANIFEST.in +├───pylintrc ├───pyproject.toml ├───README.md ├───requirements.txt -└───setup.py +└───requirements_tox.txt ``` ### Installing in Editable Mode diff --git a/README.md b/README.md index 18e6153..c4d699b 100644 --- a/README.md +++ b/README.md @@ -1,6 +1,6 @@ # QuantRL: Quantum Control using Reinforcement Learning -![Latest Version](https://img.shields.io/badge/version-0.0.8-red?style=for-the-badge) +![Latest Version](https://img.shields.io/badge/version-0.0.9-red?style=for-the-badge) > A library of modules to interface deterministic and stochastic quantum models for reinforcement learning. @@ -11,26 +11,23 @@ * Support for deterministic and stochastic linear environments. * Live visualization and learning curves. -### What's New in v0.0.7 - -* Added support for measurement noise in observations. -* Updated stochastic environment for fast Wiener processes. -* Asynchronous cache-dump to speed up environment evolution. -* Callback to save best mean reward. - -### What's New in v0.0.6 +### What's New in v0.0.x * Initialize environments with any of the three backends: NumPy, PyTorch and JAX. * Solve IVPs for the popular libraries `TorchDiffEq` and `Diffrax`. +* Asynchronous cache-dump to speed up environment evolution. +* Updated stochastic environment for fast Wiener processes. +* Added support for measurement noise in observations. +* Callback to save best mean reward. ## Installation -[QuantRL](https://github.com/sampreet/quantrl) requires `Python 3.10+`, preferably installed via the [Anaconda distribution](https://www.anaconda.com/download). +[QuantRL](https://github.com/sampreet/quantrl) requires `Python 3.12+`, preferably installed via the [Anaconda distribution](https://www.anaconda.com/download). The toolbox primarily relies on `gymnasium` (for single environments) and `stable-baselines3` (for vectorized environments). All of its dependencies can be installed using: ```bash -conda install "numpy<2.0.0" scipy matplotlib tqdm pillow pandas gymnasium stable-baselines3 +python -m pip install numpy scipy matplotlib tqdm pillow pandas gymnasium stable-baselines3 ``` Additionally, to avail the PyTorch or JAX backends, the latest version of these framework (for both CPU and GPU) should be installed (preferably in different `conda` environments) using in their official documentations: [PyTorch docs](https://pytorch.org/get-started/locally/) and [JAX docs](https://jax.readthedocs.io/en/latest/installation.html). diff --git a/pylintrc b/pylintrc index fd636f5..2c41697 100644 --- a/pylintrc +++ b/pylintrc @@ -11,10 +11,9 @@ disable=too-many-lines, redefined-builtin, # catches __name__ duplicate-code, # catches __init__ unused-argument, # catches Gym.Env methods - import-outside-toplevel, # for numerical libraries + import-outside-toplevel, # for context managers not-callable, # catches JAX JIT functions - fixme - + fixme, [BASIC] attr-naming-style=any variable-naming-style=any diff --git a/pyproject.toml b/pyproject.toml index 247c3cf..90a83bc 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -20,12 +20,13 @@ keywords = [ ] readme = "README.md" license = {file = "LICENSE"} -requires-python = ">=3.10" +requires-python = ">=3.12" dependencies = [ - "numpy<2.0.0", + "numpy", "scipy", "matplotlib", "tqdm", + "rich", "pillow", "pandas", "gymnasium", diff --git a/quantrl/__init__.py b/quantrl/__init__.py index e83f115..6ff4b2f 100644 --- a/quantrl/__init__.py +++ b/quantrl/__init__.py @@ -1,2 +1,2 @@ """Module to initialize QuantRL.""" -__version__ = "0.0.8" +__version__ = "0.0.9" diff --git a/quantrl/backends/jax.py b/quantrl/backends/jax.py index 55b8673..3504bd4 100644 --- a/quantrl/backends/jax.py +++ b/quantrl/backends/jax.py @@ -6,7 +6,7 @@ __name__ = 'quantrl.backends.jax' __authors__ = ["Sampreet Kalita"] __created__ = "2024-03-10" -__updated__ = "2024-10-13" +__updated__ = "2025-04-21" # dependencies from inspect import getfullargspec @@ -64,44 +64,44 @@ def transpose( return jnp.transpose(tensor, axes=_axes) self.jit_transpose = jax.jit( - fun=transpose, + transpose, static_argnums=(1, 2) ) self.jit_repeat = jax.jit( - fun=jnp.repeat, + jnp.repeat, static_argnums=(1, 2) ) self.jit_add = jax.jit( - fun=lambda tensor_0, tensor_1, out: jnp.add(tensor_0, tensor_1), + lambda tensor_0, tensor_1, out: jnp.add(tensor_0, tensor_1), donate_argnums=(2, ) ) self.jit_matmul = jax.jit( - fun=lambda tensor_0, tensor_1, out: jnp.matmul(tensor_0, tensor_1), + lambda tensor_0, tensor_1, out: jnp.matmul(tensor_0, tensor_1), donate_argnums=(2, ) ) self.jit_dot = jax.jit( - fun=lambda tensor_0, tensor_1, out: jnp.dot(tensor_0, tensor_1), + lambda tensor_0, tensor_1, out: jnp.dot(tensor_0, tensor_1), donate_argnums=(2, ) ) self.jit_concatenate = jax.jit( - fun=lambda tensors, axis, out: jnp.concatenate(tensors, axis), + lambda tensors, axis, out: jnp.concatenate(tensors, axis), static_argnums=(1, ), donate_argnums=(2, ) ) self.jit_stack = jax.jit( - fun=lambda tensors, axis, out: jnp.stack(tensors, axis), + lambda tensors, axis, out: jnp.stack(tensors, axis), static_argnums=(1, ), donate_argnums=(2, ) ) self.jit_update = jax.jit( - fun=lambda tensor, indices, values: tensor.at[indices].set(values), + lambda tensor, indices, values: tensor.at[indices].set(values), donate_argnums=(0, ) ) diff --git a/quantrl/envs/base.py b/quantrl/envs/base.py index f3e270e..0be705e 100644 --- a/quantrl/envs/base.py +++ b/quantrl/envs/base.py @@ -6,7 +6,7 @@ __name__ = 'quantrl.envs.base' __authors__ = ["Sampreet Kalita"] __created__ = "2023-04-25" -__updated__ = "2024-10-14" +__updated__ = "2025-05-10" # dependencies from abc import ABC, abstractmethod @@ -548,7 +548,7 @@ def plot_learning_curve(self, # initialize plotter plotter = LearningCurvePlotter( axis_args=axis_args if axis_args is not None and len(axis_args) == 4 else self.default_axis_args_learning_curve, - average_over=self.average_over + average_over=self.average_over if self.average_over < data_rewards.shape[0] else int(data_rewards.shape[0] / 2) ) # update plot plotter.add_data( diff --git a/quantrl/io.py b/quantrl/io.py index 1447be8..db9e94a 100644 --- a/quantrl/io.py +++ b/quantrl/io.py @@ -6,7 +6,7 @@ __name__ = 'quantrl.io' __authors__ = ["Sampreet Kalita"] __created__ = "2023-12-07" -__updated__ = "2024-10-14" +__updated__ = "2025-04-21" # dependencies import gc @@ -96,7 +96,7 @@ def update_cache(self, # update list if self.cache is None: - self.cache = np.zeros((self.cache_dump_interval, *data.shape), dtype=np.float_) + self.cache = np.zeros((self.cache_dump_interval, *data.shape), dtype=data.dtype) self.index += 1 self.cache[self.index % self.cache_dump_interval] = data diff --git a/quantrl/solvers/measure.py b/quantrl/solvers/measure.py index ef90e2d..830c9b1 100644 --- a/quantrl/solvers/measure.py +++ b/quantrl/solvers/measure.py @@ -24,13 +24,16 @@ __name__ = 'qom.solvers.measure' __authors__ = ["Sampreet Kalita"] __created__ = "2021-01-04" -__updated__ = "2024-10-14" +__updated__ = "2025-04-21" # dependencies from typing import Union import numpy as np +# TODO: implement backends and dtypes +# currently set to 64-bit floats + class QCMSolver(): r"""Class to solve for quantum correlation measures. @@ -87,8 +90,12 @@ def __init__(self, Modes, Corrs, params:dict, cb_update=None): Corrs=Corrs ) + # set datatypes + self.numpy_complex = self.Modes.dtype if self.Modes is not None else np.complex128 + self.numpy_real = self.Corrs.dtype if self.Corrs is not None else np.float64 + # set symplectic matrix - self.Omega_s = np.kron(np.eye(2, dtype=np.float_), np.array([[0, 1], [-1, 0]], dtype=np.float_)) + self.Omega_s = np.kron(np.eye(2, dtype=self.numpy_real), np.array([[0, 1], [-1, 0]], dtype=self.numpy_real)) # set parameters self.set_params(params) @@ -155,7 +162,7 @@ def get_measures(self): _dim = (len(self.Corrs), len(self.params['measure_codes'])) # initialize measures - Measures = np.zeros(_dim, dtype=np.float_) + Measures = np.zeros(_dim, dtype=self.numpy_real) # find measures for j in range(_dim[1]): @@ -218,7 +225,7 @@ def get_submatrices(self, pos_i:int, pos_j:int): Corrs_modes = np.concatenate((np.concatenate((As, Cs), axis=2), np.concatenate((C_Ts, Bs), axis=2)), axis=1) # # correlation matrix of the two modes (slow) - # Corrs_modes = np.array([np.block([[As[i], Cs[i]], [C_Ts[i], Bs[i]]]) for i in range(len(self.Corrs))], dtype=np.float_) + # Corrs_modes = np.array([np.block([[As[i], Cs[i]], [C_Ts[i], Bs[i]]]) for i in range(len(self.Corrs))], dtype=self.numpy_real) return Corrs_modes, As, Bs, Cs @@ -283,7 +290,7 @@ def get_correlation_Pearson(self, pos_i:int, pos_j:int): mean_jj = np.mean(self.Corrs[:, pos_j, pos_j]) # Pearson correlation coefficient as a repeated array - return np.array([mean_ij / np.sqrt(mean_ii * mean_jj)] * len(self.Corrs), dtype=np.float_) + return np.array([mean_ij / np.sqrt(mean_ii * mean_jj)] * len(self.Corrs), dtype=self.numpy_real) def get_discord_Gaussian(self, pos_i:int, pos_j:int): """Method to obtain Gaussian quantum discord values [3]_. @@ -302,10 +309,10 @@ def get_discord_Gaussian(self, pos_i:int, pos_j:int): """ # initialize values - mu_pluses = np.zeros(len(self.Corrs), dtype=np.float_) - mu_minuses = np.zeros(len(self.Corrs), dtype=np.float_) - Ws = np.zeros(len(self.Corrs), dtype=np.float_) - Discord_G = np.zeros(len(self.Corrs), dtype=np.float_) + mu_pluses = np.zeros(len(self.Corrs), dtype=self.numpy_real) + mu_minuses = np.zeros(len(self.Corrs), dtype=self.numpy_real) + Ws = np.zeros(len(self.Corrs), dtype=self.numpy_real) + Discord_G = np.zeros(len(self.Corrs), dtype=self.numpy_real) # get symplectic invariants I_1s, I_2s, I_3s, I_4s = self.get_invariants( @@ -387,7 +394,7 @@ def get_entanglement_logarithmic_negativity(self, pos_i:int, pos_j:int): eigs_min = np.min(np.abs(eigs), axis=1) # initialize entanglement - Entan_ln = np.zeros_like(eigs_min, dtype=np.float_) + Entan_ln = np.zeros_like(eigs_min, dtype=self.numpy_real) # update entanglement for i, eig_min in enumerate(eigs_min): @@ -415,7 +422,7 @@ def get_entanglement_logarithmic_negativity_2(self, pos_i:int, pos_j:int): """ # initialize values - Entan_ln = np.zeros(len(self.Corrs), dtype=np.float_) + Entan_ln = np.zeros(len(self.Corrs), dtype=self.numpy_real) # symplectic invariants I_1s, I_2s, I_3s, I_4s = self.get_invariants( @@ -669,8 +676,8 @@ def get_Wigner_distributions_single_mode(Corrs, params, cb_update=None): for val in [xs, ys]: assert val is not None and isinstance(val, (list, np.ndarray)), "Solver parameters ``'wigner_xs'`` and ``'wigner_ys'`` should be either NumPy arrays or ``list``" # handle list - xs = np.array(xs, dtype=np.float_) if isinstance(xs, list) else xs - ys = np.array(ys, dtype=np.float_) if isinstance(xs, list) else ys + xs = np.array(xs, dtype=Corrs.dtype) if isinstance(xs, list) else xs + ys = np.array(ys, dtype=Corrs.dtype) if isinstance(xs, list) else ys # extract frequently used variables show_progress = params.get('show_progress', False) @@ -684,7 +691,7 @@ def get_Wigner_distributions_single_mode(Corrs, params, cb_update=None): Vects_t = np.transpose(Vects, axes=(0, 1, 3, 2)) # initialize measures - Wigners = np.zeros((dim_c, dim_m, ys.shape[0], xs.shape[0]), dtype=np.float_) + Wigners = np.zeros((dim_c, dim_m, ys.shape[0], xs.shape[0]), dtype=Corrs.dtype) # iterate over indices for j in range(dim_m): @@ -760,8 +767,8 @@ def get_Wigner_distributions_two_mode(Corrs, params, cb_update=None): for val in [xs, ys]: assert val is not None and isinstance(val, (list, np.ndarray)), "Solver parameters ``'wigner_xs'`` and ``'wigner_ys'`` should be either NumPy arrays or ``list``" # handle list - xs = np.array(xs, dtype=np.float_) if isinstance(xs, list) else xs - ys = np.array(ys, dtype=np.float_) if isinstance(xs, list) else ys + xs = np.array(xs, dtype=Corrs.dtype) if isinstance(xs, list) else xs + ys = np.array(ys, dtype=Corrs.dtype) if isinstance(xs, list) else ys # extract frequently used variables show_progress = params.get('show_progress', False) @@ -773,13 +780,13 @@ def get_Wigner_distributions_two_mode(Corrs, params, cb_update=None): # get column vectors and row vectors _X, _Y = np.meshgrid(xs, ys) _dim = (ys.shape[0], xs.shape[0], 1, 1) - Vects_a = np.concatenate((np.reshape(_X, _dim), np.zeros(_dim, dtype=np.float_)), axis=2) if indices[0][1] == 0 else np.concatenate((np.zeros(_dim, dtype=np.float_), np.reshape(_X, _dim)), axis=2) - Vects_b = np.concatenate((np.reshape(_Y, _dim), np.zeros(_dim, dtype=np.float_)), axis=2) if indices[1][1] == 0 else np.concatenate((np.zeros(_dim, dtype=np.float_), np.reshape(_Y, _dim)), axis=2) + Vects_a = np.concatenate((np.reshape(_X, _dim), np.zeros(_dim, dtype=Corrs.dtype)), axis=2) if indices[0][1] == 0 else np.concatenate((np.zeros(_dim, dtype=Corrs.dtype), np.reshape(_X, _dim)), axis=2) + Vects_b = np.concatenate((np.reshape(_Y, _dim), np.zeros(_dim, dtype=Corrs.dtype)), axis=2) if indices[1][1] == 0 else np.concatenate((np.zeros(_dim, dtype=Corrs.dtype), np.reshape(_Y, _dim)), axis=2) Vects = np.concatenate((Vects_a, Vects_b), axis=2) Vects_t = np.transpose(Vects, axes=(0, 1, 3, 2)) # initialize measures - Wigners = np.zeros((dim_c, ys.shape[0], xs.shape[0]), dtype=np.float_) + Wigners = np.zeros((dim_c, ys.shape[0], xs.shape[0]), dtype=Corrs.dtype) # correlation matrix of the ith mode As = Corrs[:, pos_i:pos_i + 2, pos_i:pos_i + 2] @@ -846,8 +853,8 @@ def validate_Modes_Corrs(Modes=None, Corrs=None, is_modes_required:bool=False, i assert Corrs is not None if is_corrs_required else True, "Missing required parameter ``Corrs``" # handle list - Modes = np.array(Modes, dtype=np.complex_) if Modes is not None and isinstance(Modes, list) else Modes - Corrs = np.array(Corrs, dtype=np.float_) if Corrs is not None and isinstance(Corrs, list) else Corrs + Modes = np.array(Modes, dtype=np.complex128) if Modes is not None and isinstance(Modes, list) else Modes + Corrs = np.array(Corrs, dtype=np.float64) if Corrs is not None and isinstance(Corrs, list) else Corrs # validate shapes assert len(Modes.shape) == 2 if Modes is not None else True, "``Modes`` should be of shape ``(dim, num_modes)``" @@ -882,7 +889,7 @@ def validate_As_Coeffs(As=None, Coeffs=None): # validate drift matrix assert isinstance(As, Union[list, np.ndarray].__args__), "``As`` should be of type ``list`` or ``numpy.ndarray``" # convert to numpy array - As = np.array(As, dtype=np.float_) if isinstance(As, list) else As + As = np.array(As, dtype=np.float64) if isinstance(As, list) else As # validate shape assert len(As.shape) == 3 and As.shape[1] == As.shape[2], "``As`` should be of shape ``(dim_0, 2 * num_modes, 2 * num_modes)``" # if coefficients are given @@ -890,7 +897,7 @@ def validate_As_Coeffs(As=None, Coeffs=None): # validate coefficients assert isinstance(Coeffs, Union[list, np.ndarray].__args__), "``Coeffs`` should be of type ``list`` or ``numpy.ndarray``" # convert to numpy array - Coeffs = np.array(Coeffs, dtype=np.float_) if isinstance(Coeffs, list) else Coeffs + Coeffs = np.array(Coeffs, dtype=np.float64) if isinstance(Coeffs, list) else Coeffs # validate shape assert len(Coeffs.shape) == 2, "``Coeffs`` should be of shape ``(dim_0, 2 * num_modes + 1)``" diff --git a/requirements.txt b/requirements.txt index 1358fcb..4baedc3 100644 --- a/requirements.txt +++ b/requirements.txt @@ -1,7 +1,8 @@ -numpy<2.0.0 +numpy scipy matplotlib tqdm +rich pillow pandas gymnasium diff --git a/requirements_tox.txt b/requirements_tox.txt index 27cddc8..155bafc 100644 --- a/requirements_tox.txt +++ b/requirements_tox.txt @@ -1,10 +1,8 @@ -pytest -pytest-cov -pylint -numpy<2.0.0 +numpy scipy matplotlib tqdm +rich pillow pandas gymnasium @@ -12,3 +10,6 @@ stable-baselines3 torchdiffeq jax diffrax +pytest +pytest-cov +pylint From a0b592ae491351a29cdd1c76e5ec443cd75df2d3 Mon Sep 17 00:00:00 2001 From: Sampreet Kalita <9553215+Sampreet@users.noreply.github.com> Date: Mon, 12 May 2025 10:59:38 +0100 Subject: [PATCH 2/5] PyTorch GPU Support --- .github/workflows/python-tox.yml | 2 - .gitignore | 3 +- CHANGELOG.md | 19 +++++ CONTRIBUTING.md | 19 ++++- README.md | 54 ++++++++------- docs/source/conf.py | 2 +- pylintrc | 1 + pyproject.toml | 19 +++-- quantrl/__init__.py | 2 +- quantrl/backends/context_manager.py | 50 ++++++++------ quantrl/backends/jax.py | 8 +-- quantrl/backends/torch.py | 24 ++++--- quantrl/envs/base.py | 19 ++--- quantrl/envs/deterministic.py | 26 +++---- quantrl/envs/stochastic.py | 10 +-- quantrl/solvers/base.py | 4 +- quantrl/solvers/context_manager.py | 41 ++++++----- quantrl/solvers/jax.py | 10 +-- quantrl/solvers/numpy.py | 4 +- quantrl/solvers/torch.py | 6 +- requirements.txt | 2 + requirements_tox.txt | 7 +- tests/__init__.py | 0 tests/envs/__init__.py | 0 tests/envs/test_stochastic.py | 103 ++++++++++++++++++++++++++++ 25 files changed, 302 insertions(+), 133 deletions(-) create mode 100644 tests/__init__.py create mode 100644 tests/envs/__init__.py create mode 100644 tests/envs/test_stochastic.py diff --git a/.github/workflows/python-tox.yml b/.github/workflows/python-tox.yml index bdda448..52c22b0 100644 --- a/.github/workflows/python-tox.yml +++ b/.github/workflows/python-tox.yml @@ -13,8 +13,6 @@ jobs: matrix: os: [ ubuntu-latest, - macos-latest, - windows-latest, ] steps: diff --git a/.gitignore b/.gitignore index cf078ac..30374c2 100644 --- a/.gitignore +++ b/.gitignore @@ -2,9 +2,9 @@ .vscode/ # cache files -temp*/ data*/ img*/ +temp*/ # Jupyter Notebook Checkpoints *.ipynb_checkpoints/ @@ -49,6 +49,7 @@ pip-delete-this-directory.txt # Unit test / coverage reports htmlcov/ +.pytest_cache/ .tox/ .coverage .coverage.* diff --git a/CHANGELOG.md b/CHANGELOG.md index b29e6ec..c9182b7 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -1,5 +1,24 @@ # Changelog +## 2025/05/12 - 00 - v0.0.10 - PyTorch GPU Support +* Removed CI builds for MacOS and Windows. +* Changes to `quantrl.backends` package: + * Updated context managers to catch import errors. + * Renamed `jax.JaxBackend` class to `jax.JAXBackend`. + * Minor fixes to CUDA options for `torch.PyTorchBackend` class. +* Minor fixes to `quantrl.envs` package modules. +* Changes to `quantrl.solvers` package: + * Update context managers to catch import errors. + * Minor fixes to `jax` and `numpy` modules. + * Added ``'tsit5'`` option in `pytorch.TorchDiffEqIVPSolver` class. +* Bumped version to `0.0.10` in `docs/source/conf.py` and `quantrl.__init__.py`. +* Added tests for `quantrl.envs.stochastic` module. +* Updated `.gitignore` and `pylintrc`. +* Changes to `pyproject.toml`: + * Removed optional import for PyTorch, which is now installed as a dependency. + * Added `pytest` tests and renamed ``'lint'`` environment to ``'test'`` for `tox`. +* Updated `CONTRIBUTING`, `requirements` and `README`. + ## 2025/05/10 - 00 - v0.0.9 - Update Requirements * Added support for NumPy version `2.0+` with minimum Python version `3.12`. * Minor fixes to `quantrl.backends.jax` and `quantrl.io` modules. diff --git a/CONTRIBUTING.md b/CONTRIBUTING.md index 72b6e6d..c10f258 100644 --- a/CONTRIBUTING.md +++ b/CONTRIBUTING.md @@ -66,12 +66,27 @@ ROOT_DIR/ ### Installing in Editable Mode -To install the package in editable mode, execute the following from *outside* the top-level directory, `ROOT_DIR`, inside which `setup.py` is located: +To install the package in editable mode, execute the following from *inside* the top-level directory, `ROOT_DIR`, inside which `setup.py` is located: ```bash -pip install -e ROOT_DIR +python -m pip install -r requirements.txt +python -m pip install -e . ``` +To install the JAX dependencies, use: + +```bash +python -m pip install -e .[jax-cpu] +``` + +for the CPU version, or, + +```bash +python -m pip install -e .[jax-gpu] +``` + +for the GPU version with CUDA 12. + ### Building the Documentation To auto-generate and build the API documentation, navigate to the `ROOT_DIR/docs` folder and execute: diff --git a/README.md b/README.md index c4d699b..34d8896 100644 --- a/README.md +++ b/README.md @@ -1,53 +1,59 @@ # QuantRL: Quantum Control using Reinforcement Learning -![Latest Version](https://img.shields.io/badge/version-0.0.9-red?style=for-the-badge) +![Latest Version](https://img.shields.io/badge/version-0.0.10-red?style=for-the-badge) -> A library of modules to interface deterministic and stochastic quantum models for reinforcement learning. +> A backend-agnostic library of modules to interface deterministic and stochastic quantum models for reinforcement learning. ### Key Features! -* Quickly interface environments for Reinforcement Learning using Stable-Baselines3. -* Run multiple environments in parallel using vectorized inheritable classes. -* Support for deterministic and stochastic linear environments. -* Live visualization and learning curves. +* Quickly interface environments with any of the three backends: NumPy, PyTorch and JAX. +* Run multiple RL environments in parallel using vectorized inheritable classes. +* Evolve deterministic and stochastic environments with asynchronous saves. +* Visualize evolutions and plot learning curves seamlessly. -### What's New in v0.0.x +### What's New! -* Initialize environments with any of the three backends: NumPy, PyTorch and JAX. -* Solve IVPs for the popular libraries `TorchDiffEq` and `Diffrax`. -* Asynchronous cache-dump to speed up environment evolution. -* Updated stochastic environment for fast Wiener processes. -* Added support for measurement noise in observations. -* Callback to save best mean reward. +* Support for NumPy 2.x.x. +* ``'tsit5'`` solver in PyTorch. + +For a complete list of changes, see [CHANGELOG.md](CHANGELOG.md). ## Installation [QuantRL](https://github.com/sampreet/quantrl) requires `Python 3.12+`, preferably installed via the [Anaconda distribution](https://www.anaconda.com/download). -The toolbox primarily relies on `gymnasium` (for single environments) and `stable-baselines3` (for vectorized environments). -All of its dependencies can be installed using: +It's base dependencies can be installed using: ```bash -python -m pip install numpy scipy matplotlib tqdm pillow pandas gymnasium stable-baselines3 +python -m pip install numpy scipy matplotlib tqdm rich pillow pandas ``` -Additionally, to avail the PyTorch or JAX backends, the latest version of these framework (for both CPU and GPU) should be installed (preferably in different `conda` environments) using in their official documentations: [PyTorch docs](https://pytorch.org/get-started/locally/) and [JAX docs](https://jax.readthedocs.io/en/latest/installation.html). -After successful installation, the corresponding libraries (`torchdiffeq` for PyTorch and `diffrax` for JAX) can be installed using PIP as: +The default backend for the library uses vanilla NumPy and Scipy. +To avail the JAX or PyTorch backends, the latest version of these framework (for both CPU and GPU) should be installed (preferably in different `conda` environments) using in their official documentations: [JAX docs](https://jax.readthedocs.io/en/latest/installation.html) and [PyTorch docs](https://pytorch.org/get-started/locally/). +After successful installation, the corresponding libraries (`diffrax` for JAX and `torchdiffeq` for PyTorch) can be installed using PIP. + +For the CPU versions, use: ```bash -pip install torchdiffeq +python -m pip install torch torchdiffeq jax diffrax ``` -or, +For the GPU versions with CUDA 12 support, use: ```bash -pip install jax -pip install diffrax +python -m pip install torch --index-url https://download.pytorch.org/whl/cu126 +python -m pip install torchdiffeq "jax[cuda12]" diffrax ``` -To install JAX with GPU support, use `jax[cuda12]`. - ***Note: JAX-GPU support for Windows and MacOS is still limited but it runs well in WSL2.*** +QuantRL primarily relies on `gymnasium` (for single environments) and `stable-baselines3` (for vectorized environments). + +These can be installed using: + +```bash +python -m pip install gymnasium stable-baselines3 +``` + Finally, to install the latest version of `quantrl`, execute: ```bash diff --git a/docs/source/conf.py b/docs/source/conf.py index c58d55b..5379f07 100644 --- a/docs/source/conf.py +++ b/docs/source/conf.py @@ -22,7 +22,7 @@ author = 'Sampreet Kalita' # The full version, including alpha/beta/rc tags -release = '0.0.7' +release = '0.0.10' # -- General configuration --------------------------------------------------- diff --git a/pylintrc b/pylintrc index 2c41697..f6c0d48 100644 --- a/pylintrc +++ b/pylintrc @@ -13,6 +13,7 @@ disable=too-many-lines, unused-argument, # catches Gym.Env methods import-outside-toplevel, # for context managers not-callable, # catches JAX JIT functions + unnecessary-lambda, # catches JAX JIT functions fixme, [BASIC] attr-naming-style=any diff --git a/pyproject.toml b/pyproject.toml index 90a83bc..e6073c7 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -1,5 +1,5 @@ [build-system] -requires = ["cython", "setuptools>=61", "wheel"] +requires = ["cython", "setuptools", "wheel"] build-backend = "setuptools.build_meta" [project] @@ -29,6 +29,8 @@ dependencies = [ "rich", "pillow", "pandas", + "torch", + "torchdiffeq", "gymnasium", "stable-baselines3", ] @@ -50,12 +52,6 @@ jax-gpu = [ "jax[cuda12]", "diffrax", ] -torch = [ - "torch", - "torchvision", - "torchaudio", - "torchdiffeq", -] [project.urls] Homepage = "https://github.com/sampreet/quantrl" @@ -77,12 +73,13 @@ requires = tox>=4 virtualenv>=20 env_list = - lint + test -[testenv:lint] -description = run pylint under {base_python} +[testenv:test] +description = run tests under {base_python} deps = -r requirements_tox.txt commands = - pylint quantrl + pylint quantrl tests + pytest tests --cov=quantrl --cov-report=term-missing """ diff --git a/quantrl/__init__.py b/quantrl/__init__.py index 6ff4b2f..017dc0d 100644 --- a/quantrl/__init__.py +++ b/quantrl/__init__.py @@ -1,2 +1,2 @@ """Module to initialize QuantRL.""" -__version__ = "0.0.9" +__version__ = "0.0.10" diff --git a/quantrl/backends/context_manager.py b/quantrl/backends/context_manager.py index b5c5be2..0f4a217 100644 --- a/quantrl/backends/context_manager.py +++ b/quantrl/backends/context_manager.py @@ -6,48 +6,56 @@ __name__ = 'quantrl.backends.context_manager' __authors__ = ["Sampreet Kalita"] __created__ = "2024-10-09" -__updated__ = "2024-10-13" +__updated__ = "2025-05-11" # quantrl modules from .base import BaseBackend -BACKEND_INSTANCES = {} +BACKENDS = {} -# TODO: validate arguments def get_backend_instance( library:str, precision:str='double', - device:str='gpu' + device:str='cuda' ) -> BaseBackend: """Method to obtain an instantiated backend. Parameters ---------- library: str - Name of the library. Options are ``'jax'``, ``'numpy'`` and ``'torch'``. + Name of the library. Options are ``'jax'``, ``'torch'`` and ``'numpy'``. precision: str, default='double' Precision of the numerical values in the backend. Options are ``'single'`` and ``'double'``. - device: str, default='gpu' - Device for the backend. Options are ``'cpu'`` and ``'gpu'``. + device: str, default='cuda' + Device for the backend. Options are ``'cpu'`` and ``'cuda'``. Returns ------- - backend: :class:`quantrl.backends.base.BaseBackend` + Backend: :class:`quantrl.backends.base.BaseBackend` The instantiated backend. """ - if library in BACKEND_INSTANCES: - return BACKEND_INSTANCES[library] + if library in BACKENDS: + return BACKENDS[library] + if 'jax' in library.lower(): - from .jax import JaxBackend - BACKEND_INSTANCES['jax'] = JaxBackend(precision=precision) - library = 'jax' - elif 'torch' in library.lower(): + try: + from .jax import JAXBackend + BACKENDS['jax'] = JAXBackend(precision=precision) + library = 'jax' + return BACKENDS[library] + # use PyTorch if JAX is not installed + except ImportError: + print("JAX not installed, defaulting to PyTorch") + library = 'torch' + + if 'torch' in library.lower(): from .torch import TorchBackend - BACKEND_INSTANCES['torch'] = TorchBackend(precision=precision, device=device) + BACKENDS['torch'] = TorchBackend(precision=precision, device=device) library = 'torch' - else: - assert 'numpy' in library.lower(), 'parameter `library` can be either `"jax"`, `"numpy"` or `"pytorch"`' - from .numpy import NumPyBackend - BACKEND_INSTANCES['numpy'] = NumPyBackend(precision=precision) - library = 'numpy' - return BACKEND_INSTANCES[library] + return BACKENDS[library] + + assert 'numpy' in library.lower(), "parameter ``library`` can be either ``'jax'`, ``'torch'`` or ``'numpy'``" + from .numpy import NumPyBackend + BACKENDS['numpy'] = NumPyBackend(precision=precision) + library = 'numpy' + return BACKENDS[library] diff --git a/quantrl/backends/jax.py b/quantrl/backends/jax.py index 3504bd4..c040ace 100644 --- a/quantrl/backends/jax.py +++ b/quantrl/backends/jax.py @@ -6,7 +6,7 @@ __name__ = 'quantrl.backends.jax' __authors__ = ["Sampreet Kalita"] __created__ = "2024-03-10" -__updated__ = "2025-04-21" +__updated__ = "2025-05-11" # dependencies from inspect import getfullargspec @@ -21,7 +21,7 @@ # TODO: Implement buffers # TODO: Implement equinox -class JaxBackend(BaseBackend): +class JAXBackend(BaseBackend): """Backend to interface the JAX library. Refer to :class:`quantrl.backends.base.BaseBackend` for further documentation. @@ -69,7 +69,7 @@ def transpose( ) self.jit_repeat = jax.jit( - jnp.repeat, + lambda tensor, repeats, axis: jnp.repeat(tensor, repeats, axis), static_argnums=(1, 2) ) @@ -122,7 +122,7 @@ def convert_to_numpy(self, tensor, dtype:str=None ) -> np.ndarray: - return np.array(tensor, dtype=self.dtype_from_str( + return np.asarray(tensor, dtype=self.dtype_from_str( dtype=dtype, numpy=True ) if dtype is not None else None) diff --git a/quantrl/backends/torch.py b/quantrl/backends/torch.py index 863e7be..a61fca1 100644 --- a/quantrl/backends/torch.py +++ b/quantrl/backends/torch.py @@ -6,7 +6,7 @@ __name__ = 'quantrl.backends.torch' __authors__ = ["Sampreet Kalita"] __created__ = "2024-03-10" -__updated__ = "2024-10-13" +__updated__ = "2025-05-11" # dependencies import numpy as np @@ -22,12 +22,12 @@ class TorchBackend(BaseBackend): ---------- precision: str, default='double' Precision of the numerical values in the backend. Options are ``'single'`` and ``'double'``. - device: str, default='gpu' - Device for the backend. Options are ``'cpu'`` and ``'gpu'``. + device: str, default='cuda' + Device for the backend. Options are ``'cpu'`` and ``'cuda'``. """ def __init__(self, precision:str='double', - device:str='gpu' + device:str='cuda' ): # initialize BaseBackend super().__init__( @@ -38,8 +38,8 @@ def __init__(self, ) # set default device - assert 'cpu' in device or 'gpu' in device, "Invalid precision opted, options are ``'cpu'`` and ``'gpu'``." - if 'gpu' in device and not torch.cuda.is_available(): + assert 'cpu' in device or 'cuda' in device, "Invalid precision opted, options are ``'cpu'`` and ``'cuda'``." + if 'cuda' in device and not torch.cuda.is_available(): print("CUDA not available, defaulting to ``'cpu'``") device = 'cpu' torch.set_default_device(device) @@ -62,7 +62,15 @@ def convert_to_numpy(self, tensor, dtype:str=None ) -> np.ndarray: - return np.asarray(tensor.detach().cpu().numpy() if self.device == 'cuda' else tensor.numpy(), dtype=self.dtype_from_str( + if self.is_typed( + tensor=tensor, + dtype=dtype + ): + return np.asarray(tensor.detach().cpu().numpy() if self.device == 'cuda' else tensor.numpy(), dtype=self.dtype_from_str( + dtype=dtype, + numpy=True + ) if dtype is not None else None) + return np.asarray(tensor, dtype=self.dtype_from_str( dtype=dtype, numpy=True ) if dtype is not None else None) @@ -73,7 +81,7 @@ def generator(self, if self.seed_sequence is None: self.seed_sequence = self.get_seedsequence(seed) generator = torch.Generator(device=self.device) - generator.manual_seed(self.seed_sequence.spawn(1)[0]) + generator.manual_seed(int(self.seed_sequence.spawn(1)[0].generate_state(1)[0])) return generator def integers(self, diff --git a/quantrl/envs/base.py b/quantrl/envs/base.py index 0be705e..181c5e0 100644 --- a/quantrl/envs/base.py +++ b/quantrl/envs/base.py @@ -6,7 +6,7 @@ __name__ = 'quantrl.envs.base' __authors__ = ["Sampreet Kalita"] __created__ = "2023-04-25" -__updated__ = "2025-05-10" +__updated__ = "2025-05-11" # dependencies from abc import ABC, abstractmethod @@ -24,7 +24,6 @@ from ..io import FileIO from ..plotters import TrajectoryPlotter, LearningCurvePlotter -# TODO: Interface ConsoleIO # TODO: Support for different number of states and observables class BaseEnv(ABC): @@ -521,7 +520,7 @@ def plot_learning_curve(self, """ # validate arguments - assert data_rewards is not None or n_episodes is not None, "either one of the parameters ``n_episodes`` or ``data_rewards`` should be provided" + assert data_rewards is not None or n_episodes is not None, "either one of the parameters ``data_rewards`` or ``n_episodes`` should be provided" # extract frequently used variables _idx_s = self._idx_s @@ -941,8 +940,9 @@ def evolve(self, self.io.update_cache( data=self.all_data if self.cache_all_data else self.data ) - - # update plot + # update episode reward + self.data_rewards.append(self.rewards) + # update plotter if self.plot and self.traj_idx % self.plot_interval == 0: self.plotter.plot_lines( xs=self.T_norm, @@ -953,6 +953,7 @@ def evolve(self, # close environment if close: + self.reset() self.close( save=False ) @@ -1422,7 +1423,9 @@ def evolve(self, print("Batch truncated") break - # update plot + # update episode reward + self.data_rewards.append(self.rewards) + # update plotter if self.plot: for _i in tqdm( range(len(self.plotter_env_idxs)), @@ -1439,11 +1442,9 @@ def evolve(self, ) self.plotter.hold_plot() - # update episode reward - self.data_rewards.append(self.rewards) - # close environment if close: + self.reset() self.close( save=save ) diff --git a/quantrl/envs/deterministic.py b/quantrl/envs/deterministic.py index 90a4c1a..0801fca 100644 --- a/quantrl/envs/deterministic.py +++ b/quantrl/envs/deterministic.py @@ -6,7 +6,7 @@ __name__ = 'quantrl.envs.deterministic' __authors__ = ["Sampreet Kalita"] __created__ = "2023-04-25" -__updated__ = "2024-10-14" +__updated__ = "2025-05-11" # quantrl modules from ..backends.context_manager import get_backend_instance @@ -57,7 +57,7 @@ class LinearizedHOEnv(BaseGymEnv): data_idxs: list Indices of the data to store into the ``data`` attribute. The indices can be selected from the complete set of values at each point of time (total ``1 + n_actions + n_observations + n_properties + 1`` elements in the same order, where the first element is the time and the last element is the reward). backend_library: str, default='numpy' - Solver to use for each step. Options are ``'torch'`` for PyTorch-based solvers, ``'jax'`` for JAX-based solvers and ``'numpy'`` for NumPy/SciPy-based solvers. + Solver to use for each step. Options are ``'jax'`` for JAX-based solvers, ``'torch'`` for PyTorch-based solvers and ``'numpy'`` for NumPy/SciPy-based solvers. backend_precision: str, default='double' Precision of the numerical values in the backend. Options are ``'single'`` and ``'double'``. backend_device: str, default='cuda' @@ -70,7 +70,7 @@ class LinearizedHOEnv(BaseGymEnv): ============ ================================================ key value ============ ================================================ - ode_method (*str*) method used to solve the ODEs/DDEs. Available options are ``'dopri8'``, ``'dopri5'``, ``'bosh3'``, ``'fehlberg2'`` and ``'adaptive_huen'`` for a TorchDiffEq-based solver, ``'dopri8'``, ``'dopri5'`` and ``'tsit5'`` for a Diffrax-based solver and ``'BDF'``, ``'DOP853'``, ``'LSODA'``, ``'Radau'``, ``'RK23'``, ``'RK45'``, ``'dop853'``, ``'dopri5'``, ``'lsoda'``, ``'zvode'`` and ``'vode'`` for a SciPy-based solver. Default is ``'vode'``. + ode_method (*str*) method used to solve the ODEs/DDEs. Available options are ``'dopri5'``, ``'dopri8'`` and ``'tsit5'`` for a Diffrax-based solver, ``'adaptive_huen'``, ``'bosh3'``, ``'dopri5'``, ``'dopri8'``, ``'fehlberg2'`` and ``'tsit5'`` for a TorchDiffEq-based solver and ``'BDF'``, ``'DOP853'``, ``'LSODA'``, ``'Radau'``, ``'RK23'``, ``'RK45'``, ``'dop853'``, ``'dopri5'``, ``'lsoda'``, ``'vode'`` and ``'zvode'`` for a SciPy-based solver. Default is ``'dopri5'``. ode_atol (*float*) absolute tolerance of the ODE/DDE solver. Default is ``1e-9``. ode_rtol (*float*) relative tolerance of the ODE/DDE solver. Default is ``1e-6``. ============ ================================================ @@ -80,13 +80,13 @@ class LinearizedHOEnv(BaseGymEnv): """dict: Default parameters of the environment.""" default_ode_solver_params = { - 'ode_method': 'vode', + 'ode_method': 'dopri5', 'ode_atol': 1e-9, 'ode_rtol': 1e-6 } """dict: Default parameters of the ODE solver.""" - backend_libraries = ['torch', 'jax', 'numpy'] + backend_libraries = ['jax', 'torch', 'numpy'] """list: Available backend libraries.""" def __init__(self, @@ -106,14 +106,14 @@ def __init__(self, data_idxs:list, backend_library:str='numpy', backend_precision:str='double', - backend_device:str='gpu', + backend_device:str='cuda', dir_prefix:str='data', **kwargs ): """Class constructor for LinearizedHOEnv.""" # validate arguments - assert backend_library in self.backend_libraries, f"parameter ``solver_type`` should be one of ``{self.backend_libraries}``" + assert backend_library in self.backend_libraries, f"parameter ``backend_library`` should be one of ``{self.backend_libraries}``" # select backend backend = get_backend_instance( @@ -483,7 +483,7 @@ class LinearizedHOVecEnv(BaseSB3Env): data_idxs: list Indices of the data to store into the ``data`` attribute. The indices can be selected from the complete set of values at each point of time (total ``1 + n_actions + n_observations + n_properties + 1`` elements in the same order, where the first element is the time and the last element is the reward). backend_library: str, default='numpy' - Solver to use for each step. Options are ``'torch'`` for PyTorch-based solvers, ``'jax'`` for JAX-based solvers and ``'numpy'`` for NumPy/SciPy-based solvers. + Solver to use for each step. Options are ``'jax'`` for JAX-based solvers, ``'torch'`` for PyTorch-based solvers and ``'numpy'`` for NumPy/SciPy-based solvers. backend_precision: str, default='double' Precision of the numerical values in the backend. Options are ``'single'`` and ``'double'``. backend_device: str, default='cuda' @@ -496,7 +496,7 @@ class LinearizedHOVecEnv(BaseSB3Env): ============ ================================================ key value ============ ================================================ - ode_method (*str*) method used to solve the ODEs/DDEs. Available options are ``'dopri8'``, ``'dopri5'``, ``'bosh3'``, ``'fehlberg2'`` and ``'adaptive_huen'`` for a TorchDiffEq-based solver, ``'dopri8'``, ``'dopri5'`` and ``'tsit5'`` for a Diffrax-based solver and ``'BDF'``, ``'DOP853'``, ``'LSODA'``, ``'Radau'``, ``'RK23'``, ``'RK45'``, ``'dop853'``, ``'dopri5'``, ``'lsoda'``, ``'zvode'`` and ``'vode'`` for a SciPy-based solver. Default is ``'vode'``. + ode_method (*str*) method used to solve the ODEs/DDEs. Available options are ``'dopri8'``, ``'dopri5'`` and ``'tsit5'`` for a Diffrax-based solver, ``'dopri8'``, ``'dopri5'``, ``'bosh3'``, ``'fehlberg2'`` and ``'adaptive_huen'`` for a TorchDiffEq-based solver and ``'BDF'``, ``'DOP853'``, ``'LSODA'``, ``'Radau'``, ``'RK23'``, ``'RK45'``, ``'dop853'``, ``'dopri5'``, ``'lsoda'``, ``'zvode'`` and ``'vode'`` for a SciPy-based solver. Default is ``'vode'``. ode_atol (*float*) absolute tolerance of the ODE/DDE solver. Default is ``1e-9``. ode_rtol (*float*) relative tolerance of the ODE/DDE solver. Default is ``1e-6``. ============ ================================================ @@ -506,13 +506,13 @@ class LinearizedHOVecEnv(BaseSB3Env): """dict: Default parameters of the environment.""" default_ode_solver_params = { - 'ode_method': 'vode', + 'ode_method': 'dopri5', 'ode_atol': 1e-9, 'ode_rtol': 1e-6 } """dict: Default parameters of the ODE solver.""" - backend_libraries = ['torch', 'jax', 'numpy'] + backend_libraries = ['jax', 'torch', 'numpy'] """list: Available backend libraries.""" def __init__(self, @@ -533,14 +533,14 @@ def __init__(self, data_idxs:list, backend_library:str='numpy', backend_precision:str='double', - backend_device:str='gpu', + backend_device:str='cuda', dir_prefix:str='data', **kwargs ): """Class constructor for LinearizedHOEnv.""" # validate arguments - assert backend_library in self.backend_libraries, f"parameter ``solver_type`` should be one of ``{self.backend_libraries}``" + assert backend_library in self.backend_libraries, f"parameter ``backend_library`` should be one of ``{self.backend_libraries}``" # select backend backend = get_backend_instance( diff --git a/quantrl/envs/stochastic.py b/quantrl/envs/stochastic.py index bf0f8f2..42c6c0f 100644 --- a/quantrl/envs/stochastic.py +++ b/quantrl/envs/stochastic.py @@ -6,7 +6,7 @@ __name__ = 'quantrl.envs.stochastic' __authors__ = ["Sampreet Kalita"] __created__ = "2023-04-25" -__updated__ = "2024-10-14" +__updated__ = "2025-05-11" # dependencies import numpy as np @@ -57,7 +57,7 @@ class LinearEnv(BaseGymEnv): data_idxs: list Indices of the data to store into the ``data`` attribute. The indices can be selected from the complete set of values at each point of time (total ``1 + n_actions + n_observations + n_properties + 1`` elements in the same order, where the first element is the time and the last element is the reward). backend_library: str, default='numpy' - Solver to use for each step. Options are ``'torch'`` for PyTorch-based solvers, ``'jax'`` for JAX-based solvers and ``'numpy'`` for NumPy/SciPy-based solvers. + Solver to use for each step. Options are ``'jax'`` for JAX-based solvers, ``'torch'`` for PyTorch-based solvers and ``'numpy'`` for NumPy/SciPy-based solvers. backend_precision: str, default='double' Precision of the numerical values in the backend. Options are ``'single'`` and ``'double'``. backend_device: str, default='cuda' @@ -71,7 +71,7 @@ class LinearEnv(BaseGymEnv): default_params = {} """dict: Default parameters of the environment.""" - backend_libraries = ['torch', 'jax', 'numpy'] + backend_libraries = ['jax', 'torch', 'numpy'] """list: Available backend libraries.""" def __init__(self, @@ -89,14 +89,14 @@ def __init__(self, data_idxs:list, backend_library:str='numpy', backend_precision:str='double', - backend_device:str='gpu', + backend_device:str='cuda', dir_prefix:str='data', **kwargs ): """Class constructor for LinearEnv.""" # validate arguments - assert backend_library in self.backend_libraries, f"parameter ``solver_type`` should be one of ``{self.backend_libraries}``" + assert backend_library in self.backend_libraries, f"parameter ``backend_library`` should be one of ``{self.backend_libraries}``" # select backend backend = get_backend_instance( diff --git a/quantrl/solvers/base.py b/quantrl/solvers/base.py index aaad82c..464d5a5 100644 --- a/quantrl/solvers/base.py +++ b/quantrl/solvers/base.py @@ -6,12 +6,12 @@ __name__ = 'quantrl.solvers.base' __authors__ = ["Sampreet Kalita"] __created__ = "2024-03-10" -__updated__ = "2024-10-14" +__updated__ = "2025-05-12" # dependencies from abc import ABC, abstractmethod -from tqdm import tqdm +from tqdm.rich import tqdm # quantrl modules from ..backends.base import BaseBackend diff --git a/quantrl/solvers/context_manager.py b/quantrl/solvers/context_manager.py index 1dfcca1..7259b97 100644 --- a/quantrl/solvers/context_manager.py +++ b/quantrl/solvers/context_manager.py @@ -6,12 +6,12 @@ __name__ = 'quantrl.solvers.context_manager' __authors__ = ["Sampreet Kalita"] __created__ = "2024-10-09" -__updated__ = "2024-10-14" +__updated__ = "2025-05-11" # quantrl modules from .base import BaseIVPSolver -IVP_SOLVERS = {} +SOLVERS_IVP = {} def get_IVP_solver( library:str @@ -21,26 +21,35 @@ def get_IVP_solver( Parameters ---------- library: str - Name of the library. Options are ``'jax'``, ``'numpy'`` and ``'torch'``. + Name of the library. Options are ``'jax'``, ``'torch'`` and ``'numpy'``. Returns ------- IVPSolver: :class:`quantrl.solvers.base.BaseIVPSolver` The IVP solver class. """ - if library in IVP_SOLVERS: - return IVP_SOLVERS[library] + if library in SOLVERS_IVP: + return SOLVERS_IVP[library] + if 'jax' in library.lower(): - from .jax import DiffraxIVPSolver - IVP_SOLVERS['jax'] = DiffraxIVPSolver - library = 'jax' - elif 'torch' in library.lower(): + try: + from .jax import DiffraxIVPSolver + SOLVERS_IVP['jax'] = DiffraxIVPSolver + library = 'jax' + return SOLVERS_IVP[library] + # use PyTorch if JAX is not installed + except ImportError: + print("JAX not installed, defaulting to PyTorch") + library = 'torch' + + if 'torch' in library.lower(): from .torch import TorchDiffEqIVPSolver - IVP_SOLVERS['torch'] = TorchDiffEqIVPSolver + SOLVERS_IVP['torch'] = TorchDiffEqIVPSolver library = 'torch' - else: - assert 'numpy' in library.lower(), 'parameter `library` can be either `"jax"`, `"numpy"` or `"pytorch"`' - from .numpy import SciPyIVPSolver - IVP_SOLVERS['numpy'] = SciPyIVPSolver - library = 'numpy' - return IVP_SOLVERS[library] + return SOLVERS_IVP[library] + + assert 'numpy' in library.lower(), "parameter ``library`` can be either ``'jax'`, ``'torch'`` or ``'numpy'``" + from .numpy import SciPyIVPSolver + SOLVERS_IVP['numpy'] = SciPyIVPSolver + library = 'numpy' + return SOLVERS_IVP[library] diff --git a/quantrl/solvers/jax.py b/quantrl/solvers/jax.py index 504e581..8da43b6 100644 --- a/quantrl/solvers/jax.py +++ b/quantrl/solvers/jax.py @@ -6,14 +6,14 @@ __name__ = 'quantrl.solvers.jax' __authors__ = ["Sampreet Kalita"] __created__ = "2024-03-10" -__updated__ = "2024-10-14" +__updated__ = "2025-05-11" # dependencies import jax import diffrax as dfx # quantrl modules -from ..backends.jax import JaxBackend +from ..backends.jax import JAXBackend from .base import BaseIVPSolver # TODO: Implement interpolation @@ -21,7 +21,7 @@ class DiffraxIVPSolver(BaseIVPSolver): """ODE and DDE solver using Diffrax-based methods for initial-value problems. - Available methods are ``'dopri8'``, ``'dopri5'``, and ``'tsit5'``. + Available methods are ``'dopri5'``, ``'dopri8'``, and ``'tsit5'``. Refer to :class:`quantrl.backends.base.BaseIVPSolver` for its implementation. """ @@ -38,7 +38,7 @@ def __init__(self, has_delay:bool=False, func_delay=None, delay_interval:int=0, - backend:JaxBackend=None + backend:JAXBackend=None ): # initialize BaseIVPSolver super().__init__( @@ -50,7 +50,7 @@ def __init__(self, has_delay=has_delay, func_delay=jax.jit(func_delay) if func_delay is not None else None, delay_interval=delay_interval, - backend=backend if backend is not None else JaxBackend( + backend=backend if backend is not None else JAXBackend( precision='double' ) ) diff --git a/quantrl/solvers/numpy.py b/quantrl/solvers/numpy.py index a2ffeee..9e74b54 100644 --- a/quantrl/solvers/numpy.py +++ b/quantrl/solvers/numpy.py @@ -6,7 +6,7 @@ __name__ = 'quantrl.solvers.numpy' __authors__ = ["Sampreet Kalita"] __created__ = "2024-03-10" -__updated__ = "2024-10-14" +__updated__ = "2025-05-11" # dependencies import scipy.integrate as si @@ -19,7 +19,7 @@ class SciPyIVPSolver(BaseIVPSolver): """ODE and DDE solver using SciPy-based methods for initial-value problems. - Available methods are ``'BDF'``, ``'DOP853'``, ``'LSODA'``, ``'Radau'``, ``'RK23'``, ``'RK45'``, ``'dop853'``, ``'dopri5'``, ``'lsoda'``, ``'zvode'`` and ``'vode'``. + Available methods are ``'BDF'``, ``'DOP853'``, ``'LSODA'``, ``'Radau'``, ``'RK23'``, ``'RK45'``, ``'dop853'``, ``'dopri5'``, ``'lsoda'``, ``'vode'`` and ``'zvode'``. Refer to :class:`quantrl.backends.base.BaseIVPSolver` for its implementation. """ diff --git a/quantrl/solvers/torch.py b/quantrl/solvers/torch.py index 06b2f51..8fd86d0 100644 --- a/quantrl/solvers/torch.py +++ b/quantrl/solvers/torch.py @@ -6,7 +6,7 @@ __name__ = 'quantrl.solvers.torch' __authors__ = ["Sampreet Kalita"] __created__ = "2024-03-10" -__updated__ = "2024-10-14" +__updated__ = "2025-05-11" # dependencies from torchdiffeq import odeint @@ -20,12 +20,12 @@ class TorchDiffEqIVPSolver(BaseIVPSolver): """ODE and DDE solver using TorchDiffEq-based methods for initial-value problems. - Available methods are ``'dopri8'``, ``'dopri5'``, ``'bosh3'``, ``'fehlberg2'`` and ``'adaptive_huen'``. + Available methods are ``'adaptive_huen'``, ``'bosh3'``, ``'dopri5'``, ``'dopri8'````'fehlberg2'`` and ``'tsit5'``. Refer to :class:`quantrl.backends.base.BaseIVPSolver` for its implementation. """ # attributes - solver_methods = ['dopri8', 'dopri5', 'bosh3', 'fehlberg2', 'adaptive_huen'] + solver_methods = ['adaptive_huen', 'bosh3', 'dopri5', 'dopri8', 'fehlberg2', 'tsit5'] """list: TorchDiffEq-based methods.""" def __init__(self, diff --git a/requirements.txt b/requirements.txt index 4baedc3..241afd5 100644 --- a/requirements.txt +++ b/requirements.txt @@ -5,5 +5,7 @@ tqdm rich pillow pandas +torch +torchdiffeq gymnasium stable-baselines3 diff --git a/requirements_tox.txt b/requirements_tox.txt index 155bafc..09c7e62 100644 --- a/requirements_tox.txt +++ b/requirements_tox.txt @@ -5,11 +5,12 @@ tqdm rich pillow pandas -gymnasium -stable-baselines3 +torch torchdiffeq jax diffrax +gymnasium +stable-baselines3 +pylint pytest pytest-cov -pylint diff --git a/tests/__init__.py b/tests/__init__.py new file mode 100644 index 0000000..e69de29 diff --git a/tests/envs/__init__.py b/tests/envs/__init__.py new file mode 100644 index 0000000..e69de29 diff --git a/tests/envs/test_stochastic.py b/tests/envs/test_stochastic.py new file mode 100644 index 0000000..1fbabe9 --- /dev/null +++ b/tests/envs/test_stochastic.py @@ -0,0 +1,103 @@ +"""Module to test `quantrl.envs.stochastic`""" + +# dependencies +import numpy as np +import pytest + +# quantrl modules +from quantrl.envs.stochastic import LinearEnv + +class SHOEnv00(LinearEnv): + """Class to simulate a simple harmonic oscillator""" + default_params = { + 'n_th': 1e4, + } + def __init__(self, + backend_library='numpy', + seed=None, + ): + # initialize Gym environment + super().__init__( + name='SHOEnv', + desc="Simple Harmonic Oscillator Environment", + params={}, + t_norm_max=10.0, + t_norm_ssz=0.001, + t_norm_mul=2.0 * np.pi, + n_observations=2, + n_properties=0, + n_actions=1, + action_maximums=[0.0], + action_interval=100, + data_idxs=[2, 3], + observation_stds=[0.1] * 2, + backend_library=backend_library, + action_space_range=[-1.0, 1.0], + observation_space_range=[-1e12, 1e12], + seed=seed, + cache_dump_interval=100, + average_over=100, + plot=False, + ) + + # set parameters + self.Omega_norm = 1.0 + self.n_th = self.params['n_th'] + + # update drift matrix + self.A = self.backend.update( + tensor=self.A, + indices=( + [0, 1], + [1, 0], + ), + values=self.backend.convert_to_typed( + tensor=[ + self.Omega_norm, + - self.Omega_norm + ], + dtype='real', + ) + ) + + # set noise prefixes + self.noise_prefixes = self.backend.zeros( + shape=(2, ), + dtype='real' + ) + + def reset_states(self): + # set initial values of position and momentum + # with mean 0 and standard deviation n_th + 0.5 + states_0 = self.backend.convert_to_typed( + tensor=[-58.40454235, 256.72449014], + dtype='real' + ) + + return states_0 + + def get_A(self, t_idx, args): + # update drift matrix + return self.A + + def get_noise_prefixes(self, t_idx, args): + # return noise prefixes + return self.noise_prefixes + + def get_reward(self): + # thermal occupancies + ns = 0.5 * (self.Observations[:, 0]**2 + self.Observations[:, 1]**2) + self.Reward = 1.0 / ns + + return self.Reward + +@pytest.mark.parametrize( + 'backend_library', + ['jax', 'torch', 'numpy'], +) +def test_sho_env( + backend_library, + ): + """Function to test evolution.""" + env = SHOEnv00(backend_library=backend_library, seed=1234) + env.evolve(show_progress=False) From c972e3f71fad91f9df1fd0f06ea3e5fc63e5ba13 Mon Sep 17 00:00:00 2001 From: Sampreet Kalita <9553215+Sampreet@users.noreply.github.com> Date: Thu, 21 Aug 2025 08:55:53 +0100 Subject: [PATCH 3/5] Code Cleanup --- .gitignore | 1 + CHANGELOG.md | 4 + README.md | 32 +- pylintrc | 2 +- quantrl/backends/base.py | 636 +++++++++++++-------- quantrl/backends/context_manager.py | 46 +- quantrl/backends/jax.py | 302 ++++++---- quantrl/backends/numpy.py | 275 +++++---- quantrl/backends/torch.py | 306 ++++++---- quantrl/envs/base.py | 830 ++++++++++++++++++---------- quantrl/envs/deterministic.py | 570 +++++++++++-------- quantrl/envs/stochastic.py | 176 +++--- quantrl/io.py | 157 ++++-- quantrl/plotters.py | 180 ++++-- quantrl/solvers/base.py | 134 +++-- quantrl/solvers/context_manager.py | 16 +- quantrl/solvers/jax.py | 64 ++- quantrl/solvers/numpy.py | 113 ++-- quantrl/solvers/torch.py | 72 ++- quantrl/utils.py | 74 ++- 20 files changed, 2517 insertions(+), 1473 deletions(-) diff --git a/.gitignore b/.gitignore index 30374c2..154b333 100644 --- a/.gitignore +++ b/.gitignore @@ -58,6 +58,7 @@ nosetests.xml coverage.xml *,cover .hypothesis/ +.benchmarks/ # Translations *.mo diff --git a/CHANGELOG.md b/CHANGELOG.md index c9182b7..d3efa1c 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -1,5 +1,9 @@ # Changelog +## 2025/08/20 - 00 - v0.0.10 - Code Cleanup +* Minor changes to all modules except `quantrl.solvers.measure`. +* Updated indentations and line lengths of modules and `README`. + ## 2025/05/12 - 00 - v0.0.10 - PyTorch GPU Support * Removed CI builds for MacOS and Windows. * Changes to `quantrl.backends` package: diff --git a/README.md b/README.md index 34d8896..ea25c37 100644 --- a/README.md +++ b/README.md @@ -2,13 +2,17 @@ ![Latest Version](https://img.shields.io/badge/version-0.0.10-red?style=for-the-badge) -> A backend-agnostic library of modules to interface deterministic and stochastic quantum models for reinforcement learning. +> A backend-agnostic library of modules to interface +deterministic and stochastic quantum models for reinforcement learning. ### Key Features! -* Quickly interface environments with any of the three backends: NumPy, PyTorch and JAX. -* Run multiple RL environments in parallel using vectorized inheritable classes. -* Evolve deterministic and stochastic environments with asynchronous saves. +* Quickly interface environments with any of the three backends: +NumPy, PyTorch and JAX. +* Run multiple RL environments in parallel using +vectorized inheritable classes. +* Evolve deterministic and stochastic environments +with asynchronous saves. * Visualize evolutions and plot learning curves seamlessly. ### What's New! @@ -20,7 +24,9 @@ For a complete list of changes, see [CHANGELOG.md](CHANGELOG.md). ## Installation -[QuantRL](https://github.com/sampreet/quantrl) requires `Python 3.12+`, preferably installed via the [Anaconda distribution](https://www.anaconda.com/download). +[QuantRL](https://github.com/sampreet/quantrl) requires `Python 3.12+`, +preferably installed via the +[Anaconda distribution](https://www.anaconda.com/download). It's base dependencies can be installed using: ```bash @@ -28,8 +34,14 @@ python -m pip install numpy scipy matplotlib tqdm rich pillow pandas ``` The default backend for the library uses vanilla NumPy and Scipy. -To avail the JAX or PyTorch backends, the latest version of these framework (for both CPU and GPU) should be installed (preferably in different `conda` environments) using in their official documentations: [JAX docs](https://jax.readthedocs.io/en/latest/installation.html) and [PyTorch docs](https://pytorch.org/get-started/locally/). -After successful installation, the corresponding libraries (`diffrax` for JAX and `torchdiffeq` for PyTorch) can be installed using PIP. +To avail the JAX or PyTorch backends, the latest version +of these framework (CPU or GPU) should be installed +(preferably in different `conda` environments) +using in their official documentations: +[JAX docs](https://jax.readthedocs.io/en/latest/installation.html) and +[PyTorch docs](https://pytorch.org/get-started/locally/). +After successful installation, the corresponding libraries +(`diffrax` for JAX and `torchdiffeq` for PyTorch) can be installed using PIP. For the CPU versions, use: @@ -44,9 +56,11 @@ python -m pip install torch --index-url https://download.pytorch.org/whl/cu126 python -m pip install torchdiffeq "jax[cuda12]" diffrax ``` -***Note: JAX-GPU support for Windows and MacOS is still limited but it runs well in WSL2.*** +***Note: JAX-GPU support for Windows and MacOS +is still limited but it runs well in WSL2.*** -QuantRL primarily relies on `gymnasium` (for single environments) and `stable-baselines3` (for vectorized environments). +QuantRL primarily relies on `gymnasium` (for single environments) +and `stable-baselines3` (for vectorized environments). These can be installed using: diff --git a/pylintrc b/pylintrc index f6c0d48..f3e9d11 100644 --- a/pylintrc +++ b/pylintrc @@ -7,7 +7,7 @@ disable=too-many-lines, too-many-public-methods, too-many-instance-attributes, too-many-positional-arguments, - line-too-long, # no bounded lines + line-too-long, # catches docstring tables redefined-builtin, # catches __name__ duplicate-code, # catches __init__ unused-argument, # catches Gym.Env methods diff --git a/quantrl/backends/base.py b/quantrl/backends/base.py index 8fed94f..97f37c8 100644 --- a/quantrl/backends/base.py +++ b/quantrl/backends/base.py @@ -6,10 +6,11 @@ __name__ = 'quantrl.backends.base' __authors__ = ["Sampreet Kalita"] __created__ = "2024-03-10" -__updated__ = "2024-10-13" +__updated__ = "2025-08-20" # dependencies from abc import ABC, abstractmethod +from typing import Type import numpy as np @@ -22,58 +23,61 @@ class BaseBackend(ABC): Name of the backend. library: Any Numerical library used by the backend. - tensor_type: Any + dtype_tensor: Any Tensor type for the backend. precision: str, default='double' - Precision of the numerical values in the backend. Options are ``'single'`` and ``'double'``. + Precision of the numerical values in the backend. + Options are ``'single'`` and ``'double'``. """ def __init__( - self, - name:str='numpy', - library:np=np, - tensor_type:np.ndarray=np.ndarray, - precision:str='double' + self, + name:str='numpy', + library:np=np, + dtype_tensor:Type[np.ndarray]=np.ndarray, + precision:str='double', ): # validate parameters - assert precision in ['single', 'double'], "parameter ``precision`` can be either ``'single'`` or ``'double'``." + assert precision in ['single', 'double'], \ + "parameter ``precision`` can be either ``'single'`` or ``'double'``." # set attributes - self.name = name - self.library = library - self.tensor_type = tensor_type - self.precision = precision - self.dtypes = { + self.name:str = name + self.library:np = library + self.dtype_tensor:Type[np.ndarray] = dtype_tensor + self.precision:str = precision + self.dtypes:dict = { 'typed': { 'single': { 'integer': self.library.int32, 'real': self.library.float32, - 'complex': self.library.complex64 + 'complex': self.library.complex64, }, 'double': { 'integer': self.library.int64, 'real': self.library.float64, - 'complex': self.library.complex128 + 'complex': self.library.complex128, } }, 'numpy': { 'single': { 'integer': np.int32, 'real': np.float32, - 'complex': np.complex64 + 'complex': np.complex64, }, 'double': { 'integer': np.int64, 'real': np.float64, - 'complex': np.complex128 + 'complex': np.complex128, } } } - self.seed_sequence = None + self.seed_sequence:np.random.SeedSequence = None - def is_typed(self, - tensor, - dtype:str=None + def is_typed( + self, + tensor, + dtype:str=None, ) -> bool: """Method to check if a tensor is a typed tensor of given dtype. @@ -82,7 +86,9 @@ def is_typed(self, tensor: Any Given tensor. dtype: str, default=None - Broad data-type. Options are ``'integer'``, ``'real'`` and ``'complex'``. If ``None``, the data-type is not checked. + Broad data-type. + Options are ``'integer'``, ``'real'`` and ``'complex'``. + If ``None``, the data-type is not checked. Returns ------- @@ -91,15 +97,18 @@ def is_typed(self, """ _dtype = self.dtype_from_str( - dtype=dtype + dtype=dtype, ) - if isinstance(tensor, self.tensor_type): - if dtype is None or (dtype is not None and tensor.dtype == _dtype): + if isinstance(tensor, self.dtype_tensor): + if dtype is None or ( + dtype is not None and tensor.dtype == _dtype + ): return True return False - def get_seedsequence(self, - seed:int=None + def get_seedsequence( + self, + seed:int=None, ) -> np.random.SeedSequence: """Method to obtain a SeedSequence object. @@ -116,22 +125,30 @@ def get_seedsequence(self, if seed is None: entropy = np.random.randint(1234567890) else: - entropy = np.random.default_rng(seed).integers(0, 1234567890, (1, ))[0] + entropy = np.random.default_rng(seed).integers( + 0, + 1234567890, + (1, ), + )[0] return np.random.SeedSequence(entropy) @abstractmethod - def convert_to_typed(self, - tensor, - dtype:str=None + def convert_to_typed( + self, + tensor, + dtype:str=None, ): - """Method to obtain a typed tensor with given data-type from a numpy array or another typed tensor. + """Method to obtain a typed tensor with given data-type + from a numpy array or another typed tensor. Parameters ---------- tensor: Any Given tensor. dtype: str, default=None - Broad data-type. Options are ``'integer'``, ``'real'`` and ``'complex'``. If ``None``, the data-type of the original array is returned. + Broad data-type. + Options are ``'integer'``, ``'real'`` and ``'complex'``. + If ``None``, the data-type of the original array is used. Returns ------- @@ -142,9 +159,10 @@ def convert_to_typed(self, raise NotImplementedError @abstractmethod - def convert_to_numpy(self, - tensor, - dtype:str=None + def convert_to_numpy( + self, + tensor, + dtype:str=None, ) -> np.ndarray: """Method to obtain a NumPy array from a typed tensor. @@ -153,7 +171,9 @@ def convert_to_numpy(self, tensor: Any Given typed tensor. dtype: str, default=None - Broad data-type. Options are ``'integer'``, ``'real'`` and ``'complex'``. If ``None``, the data-type of the original tensor is returned. + Broad data-type. + Options are ``'integer'``, ``'real'`` and ``'complex'``. + If ``None``, the data-type of the original tensor is used. Returns ------- @@ -164,15 +184,17 @@ def convert_to_numpy(self, raise NotImplementedError @abstractmethod - def generator(self, - seed:int=None + def generator( + self, + seed:int=None, ): """Method to obtain a pseudo random number generator. Parameters ---------- seed: Any, default=None - Seed for the PRNG. If ``None``, a random seed is selected in ``[0, 1000)``. + Seed for the PRNG. + If ``None``, a random seed is selected in ``[0, 1000)``. Returns ------- @@ -183,14 +205,16 @@ def generator(self, raise NotImplementedError @abstractmethod - def integers(self, - generator, - shape:tuple, - low:int=0, - high:int=1000, - dtype:str=None + def integers( + self, + generator, + shape:tuple, + low:int=0, + high:int=1000, + dtype:str=None, ): - """Method to obtain a typed tensor containing samples from a uniform distribution in the interval ``[low, high)``. + """Method to obtain a typed tensor containing samples + from a uniform distribution in the interval ``[low, high)``. Parameters ---------- @@ -203,7 +227,9 @@ def integers(self, high: int, default=1000 Highest value (exclusive). dtype: str, default=None - Broad data-type. Options are ``'integer'``, ``'real'`` and ``'complex'``. the data-type is casted to real. + Broad data-type. + Options are ``'integer'``, ``'real'`` and ``'complex'``. + If ``None``, the data is cast to real. Returns ------- @@ -214,14 +240,16 @@ def integers(self, raise NotImplementedError @abstractmethod - def normal(self, - generator, - shape:tuple, - mean:float=0.0, - std:float=1.0, - dtype:str=None + def normal( + self, + generator, + shape:tuple, + mean:float=0.0, + std:float=1.0, + dtype:str=None, ): - """Method to obtain a typed tensor containing samples from a normal distribution. + """Method to obtain a typed tensor containing samples + from a normal distribution. Parameters ---------- @@ -234,7 +262,9 @@ def normal(self, std: float, default=1.0 Standard deviation of the distribution. dtype: str, default=None - Broad data-type. Options are ``'integer'``, ``'real'`` and ``'complex'``. the data-type is casted to real. + Broad data-type. + Options are ``'integer'``, ``'real'`` and ``'complex'``. + If ``None``, the data is cast to real. Returns ------- @@ -245,14 +275,16 @@ def normal(self, raise NotImplementedError @abstractmethod - def uniform(self, - generator, - shape:tuple, - low:float=0.0, - high:float=1.0, - dtype:str=None + def uniform( + self, + generator, + shape:tuple, + low:float=0.0, + high:float=1.0, + dtype:str=None, ): - """Method to obtain a typed tensor containing samples from a uniform distribution in the half-open interval ``[0, 1)``. + """Method to obtain a typed tensor containing samples + from a uniform distribution in the half-open interval ``[0, 1)``. Parameters ---------- @@ -265,7 +297,9 @@ def uniform(self, high: float, default=1.0 Highest value (exclusive). dtype: str, default=None - Broad data-type. Options are ``'integer'``, ``'real'`` and ``'complex'``. the data-type is casted to real. + Broad data-type. + Options are ``'integer'``, ``'real'`` and ``'complex'``. + If ``None``, the data is cast to real. Returns ------- @@ -276,10 +310,11 @@ def uniform(self, raise NotImplementedError @abstractmethod - def transpose(self, - tensor, - axis_0:int=None, - axis_1:int=None + def transpose( + self, + tensor, + axis_0:int=None, + axis_1:int=None, ): """Method to transpose a typed tensor about two axes. @@ -301,10 +336,11 @@ def transpose(self, raise NotImplementedError @abstractmethod - def repeat(self, - tensor, - repeats:int, - axis:int + def repeat( + self, + tensor, + repeats:int, + axis:int, ): """Method to repeat a typed tensor about a given axis. @@ -326,10 +362,11 @@ def repeat(self, raise NotImplementedError @abstractmethod - def add(self, - tensor_0, - tensor_1, - out + def add( + self, + tensor_0, + tensor_1, + out, ): """Method to add two typed tensors. @@ -351,12 +388,14 @@ def add(self, raise NotImplementedError @abstractmethod - def matmul(self, - tensor_0, - tensor_1, - out + def matmul( + self, + tensor_0, + tensor_1, + out, ): - """Method to obtain the matrix multiplication two typed tensors along the last two axes. + """Method to obtain the matrix multiplication + of two typed tensors along the last two axes. Parameters ---------- @@ -376,10 +415,11 @@ def matmul(self, raise NotImplementedError @abstractmethod - def dot(self, - tensor_0, - tensor_1, - out + def dot( + self, + tensor_0, + tensor_1, + out, ): """Method to obtain the dot product of two typed tensors. @@ -401,9 +441,10 @@ def dot(self, raise NotImplementedError @abstractmethod - def norm(self, - tensor, - axis + def norm( + self, + tensor, + axis, ): """Method to obtain the norm of a typed tensor along a given axis. @@ -423,10 +464,11 @@ def norm(self, raise NotImplementedError @abstractmethod - def concatenate(self, - tensors:tuple, - axis, - out + def concatenate( + self, + tensors:tuple, + axis, + out, ): """Method to concatenate multiple typed tensors along a given axis. @@ -448,10 +490,11 @@ def concatenate(self, raise NotImplementedError @abstractmethod - def stack(self, - tensors:tuple, - axis:int, - out + def stack( + self, + tensors:tuple, + axis:int, + out, ): """Method to stack multiple typed tensors along a given axis. @@ -471,12 +514,14 @@ def stack(self, return NotImplementedError @abstractmethod - def update(self, - tensor, - indices, - values + def update( + self, + tensor, + indices, + values, ): - """Method to update selected indices of a typed tensor with given values. + """Method to update the values at certain indices + of a typed tensor with given values. Parameters ---------- @@ -496,11 +541,12 @@ def update(self, raise NotImplementedError @abstractmethod - def if_else(self, - condition, - func_true, - func_false, - args + def if_else( + self, + condition, + func_true, + func_false, + args, ): """Method to execute conditional statements. @@ -524,22 +570,27 @@ def if_else(self, raise NotImplementedError @abstractmethod - def iterate_i(self, - func, - iterations_i:int, - Y, - args:tuple=None + def iterate_i( + self, + func, + iterations_i:int, + Y, + args:tuple=None, ): """Method to iterate over a single variable. Parameters ---------- func: callable - Function to iterate formatted as ``func(i, *args, *kwargs)``, where i is the index of the iteration. + Function to iterate formatted as ``func(i, *args, *kwargs)``, + where i is the index of the iteration. iterations_i: int - Number of iterations in the first variable. This results in interation indices in the open interval ``[0, iterations_i)``. + Number of iterations in the first variable. + This results in iteration indices + in the open interval ``[0, iterations_i)``. Y: Any - The tensor which is updated at each iteration, with ``Y[0]`` containing the initial values. + The tensor which is updated at each iteration, + with ``Y[0]`` containing the initial values. args: tuple Arguments for the iteration. @@ -551,16 +602,19 @@ def iterate_i(self, raise NotImplementedError - def dtype_from_str(self, - dtype:str=None, - numpy:bool=False + def dtype_from_str( + self, + dtype:str=None, + numpy:bool=False, ): """Method to obtain the data-type from a string. Parameters ---------- dtype: str, default=None - Broad data-type. Options are ``'integer'``, ``'real'`` and ``'complex'``. If ``None``, the data-type is casted to real. + Broad data-type. + Options are ``'integer'``, ``'real'`` and ``'complex'``. + If ``None``, the data is cast to real. numpy: bool, default=False Option to use NumPy data-types. @@ -569,6 +623,7 @@ def dtype_from_str(self, dtype: type Selected data-type. """ + # default dtype is the real data-type if dtype is None or dtype not in ['integer', 'real', 'complex']: dtype = 'real' @@ -606,9 +661,10 @@ def jit_update(self, tensor, indices, values): """Method to JIT-compile updation.""" return self.update(tensor, indices, values) - def empty(self, - shape:tuple, - dtype:str=None + def empty( + self, + shape:tuple, + dtype:str=None, ): """Method to create an empty typed tensor. @@ -617,7 +673,9 @@ def empty(self, shape: tuple Shape of the tensor. dtype: str, default=None - Broad data-type. Options are ``'integer'``, ``'real'`` and ``'complex'``. If ``None``, the data-type is casted to real. + Broad data-type. + Options are ``'integer'``, ``'real'`` and ``'complex'``. + If ``None``, the data is cast to real. Returns ------- @@ -625,13 +683,17 @@ def empty(self, Empty typed tensor. """ - return self.library.empty(shape, dtype=self.dtype_from_str( - dtype=dtype - )) + return self.library.empty( + shape, + dtype=self.dtype_from_str( + dtype=dtype, + ), + ) - def zeros(self, - shape:tuple, - dtype:str=None + def zeros( + self, + shape:tuple, + dtype:str=None, ): """Method to create a typed tensor of zeros. @@ -640,7 +702,9 @@ def zeros(self, shape: tuple Shape of the tensor. dtype: str, default=None - Broad data-type. Options are ``'integer'``, ``'real'`` and ``'complex'``. If ``None``, the data-type is casted to real. + Broad data-type. + Options are ``'integer'``, ``'real'`` and ``'complex'``. + If ``None``, the data is cast to real. Returns ------- @@ -648,13 +712,17 @@ def zeros(self, Typed tensor of zeros. """ - return self.library.zeros(shape, dtype=self.dtype_from_str( - dtype=dtype - )) + return self.library.zeros( + shape, + dtype=self.dtype_from_str( + dtype=dtype, + ), + ) - def ones(self, - shape:tuple, - dtype:str=None + def ones( + self, + shape:tuple, + dtype:str=None, ): """Method to create a typed tensor of ones. @@ -663,7 +731,9 @@ def ones(self, shape: tuple Shape of the tensor. dtype: str, default=None - Broad data-type. Options are ``'integer'``, ``'real'`` and ``'complex'``. If ``None``, the data-type is casted to real. + Broad data-type. + Options are ``'integer'``, ``'real'`` and ``'complex'``. + If ``None``, the data is cast to real. Returns ------- @@ -671,14 +741,18 @@ def ones(self, Typed tensor of ones. """ - return self.library.ones(shape, dtype=self.dtype_from_str( - dtype=dtype - )) + return self.library.ones( + shape, + dtype=self.dtype_from_str( + dtype=dtype, + ), + ) - def eye(self, - N:int, - M:int=None, - dtype:str=None + def eye( + self, + N:int, + M:int=None, + dtype:str=None, ): """Method to create an typed identity matrix. @@ -687,9 +761,12 @@ def eye(self, N: tuple Number of rows. M: int, defualt=None - Number of columns. if ``None``, this value is set equal to the number of rows. + Number of columns. + If ``None``, this value is set equal to the number of rows. dtype: str, default=None - Broad data-type. Options are ``'integer'``, ``'real'`` and ``'complex'``. If ``None``, the data-type is casted to real. + Broad data-type. + Options are ``'integer'``, ``'real'`` and ``'complex'``. + If ``None``, the data is cast to real. Returns ------- @@ -697,13 +774,18 @@ def eye(self, Typed identity matrix. """ - return self.library.eye(N, (M if M is not None else N), dtype=self.dtype_from_str( - dtype=dtype - )) + return self.library.eye( + N, + (M if M is not None else N), + dtype=self.dtype_from_str( + dtype=dtype, + ), + ) - def diag(self, - tensor, - dtype:str=None + def diag( + self, + tensor, + dtype:str=None, ): """Method to create an typed diagonal matrix. @@ -712,7 +794,9 @@ def diag(self, tensor: tuple Elements of the diagonal. dtype: str, default=None - Broad data-type. Options are ``'integer'``, ``'real'`` and ``'complex'``. If ``None``, the data-type of the original tensor is returned. + Broad data-type. + Options are ``'integer'``, ``'real'`` and ``'complex'``. + If ``None``, the data-type of the original tensor is used. Returns ------- @@ -720,18 +804,22 @@ def diag(self, Typed diagonal matrix. """ - return self.library.diag(self.convert_to_typed( - tensor=tensor, - dtype=dtype - )) - - def arange(self, - start:float, - stop:float, - ssz:float, - dtype:str=None + return self.library.diag( + self.convert_to_typed( + tensor=tensor, + dtype=dtype, + ), + ) + + def arange( + self, + start:float, + stop:float, + ssz:float, + dtype:str=None, ): - """Method to create a typed tensor of evenly-stepped values from ``start`` (included) to ``stop`` (excluded). + """Method to create a typed tensor of evenly-stepped values + from ``start`` (included) to ``stop`` (excluded). Parameters ---------- @@ -742,7 +830,9 @@ def arange(self, ssz: float Size of the steps. dtype: str, default=None - Broad data-type. Options are ``'integer'``, ``'real'`` and ``'complex'``. If ``None``, the data-type is casted to integer. + Broad data-type. + Options are ``'integer'``, ``'real'`` and ``'complex'``. + If ``None``, the data is cast to integer. Returns ------- @@ -750,17 +840,24 @@ def arange(self, Typed tensor of evenly-stepped values. """ - return self.library.arange(start, stop, ssz, dtype=self.dtype_from_str( - dtype=dtype if dtype is not None else 'integer' - )) + return self.library.arange( + start, + stop, + ssz, + dtype=self.dtype_from_str( + dtype=dtype if dtype is not None else 'integer', + ), + ) - def linspace(self, - start:float, - stop:float, - dim:int, - dtype:str=None + def linspace( + self, + start:float, + stop:float, + dim:int, + dtype:str=None, ): - """Method to create a typed tensor of linearly-spaced values from ``start`` to ``stop``, both inclusive. + """Method to create a typed tensor of linearly-spaced values + from ``start`` to ``stop``, both inclusive. Parameters ---------- @@ -771,7 +868,9 @@ def linspace(self, dim: int Dimension of the tensor. dtype: str, default=None - Broad data-type. Options are ``'integer'``, ``'real'`` and ``'complex'``. If ``None``, the data-type is casted to real. + Broad data-type. + Options are ``'integer'``, ``'real'`` and ``'complex'``. + If ``None``, the data is cast to real. Returns ------- @@ -779,12 +878,18 @@ def linspace(self, Typed tensor of linearly-spaced values. """ - return self.library.linspace(start, stop, dim, dtype=self.dtype_from_str( - dtype=dtype - )) + return self.library.linspace( + start, + stop, + dim, + dtype=self.dtype_from_str( + dtype=dtype, + ), + ) - def shape(self, - tensor + def shape( + self, + tensor, ) -> tuple: """Method to obtain the shape of a tensor. @@ -800,12 +905,13 @@ def shape(self, """ return tuple(self.convert_to_typed( - tensor=tensor + tensor=tensor, ).shape) - def reshape(self, - tensor, - shape:tuple + def reshape( + self, + tensor, + shape:tuple, ): """Method to reshape a typed tensor. @@ -823,11 +929,12 @@ def reshape(self, """ return self.convert_to_typed( - tensor=tensor + tensor=tensor, ).reshape(shape) - def flatten(self, - tensor + def flatten( + self, + tensor, ): """Method to flatten typed tensor. @@ -843,11 +950,12 @@ def flatten(self, """ return self.convert_to_typed( - tensor=tensor + tensor=tensor, ).flatten() - def real(self, - tensor + def real( + self, + tensor, ): """Method to obtain the real components of a complex typed tensor. @@ -862,12 +970,15 @@ def real(self, Real components of the complex typed tensor. """ - return self.library.real(self.convert_to_typed( - tensor=tensor - )) + return self.library.real( + self.convert_to_typed( + tensor=tensor, + ), + ) - def imag(self, - tensor + def imag( + self, + tensor, ): """Method to obtain the imaginary components of a complex typed tensor. @@ -882,12 +993,15 @@ def imag(self, Imaginary components of the complex typed tensor. """ - return self.library.imag(self.convert_to_typed( - tensor=tensor - )) + return self.library.imag( + self.convert_to_typed( + tensor=tensor, + ), + ) - def sqrt(self, - tensor + def sqrt( + self, + tensor, ): """Method to obtain the square root of a typed tensor. @@ -902,15 +1016,19 @@ def sqrt(self, Square root of the typed tensor. """ - return self.library.sqrt(self.convert_to_typed( - tensor=tensor - )) + return self.library.sqrt( + self.convert_to_typed( + tensor=tensor, + ), + ) - def sum(self, - tensor, - axis + def sum( + self, + tensor, + axis, ): - """Method to obtain the sum of a typed tensor along a given axis. + """Method to obtain the sum + of a typed tensor along a given axis. Parameters ---------- @@ -925,15 +1043,20 @@ def sum(self, Sum of the typed tensor along the axis. """ - return self.library.sum(self.convert_to_typed( - tensor=tensor - ), axis) + return self.library.sum( + self.convert_to_typed( + tensor=tensor, + ), + axis, + ) - def cumsum(self, - tensor, - axis + def cumsum( + self, + tensor, + axis, ): - """Method to obtain the cumulative sum of a typed tensor along a given axis. + """Method to obtain the cumulative sum + of a typed tensor along a given axis. Parameters ---------- @@ -948,12 +1071,16 @@ def cumsum(self, Cumulative um of the typed tensor along the axis. """ - return self.library.cumsum(self.convert_to_typed( - tensor=tensor - ), axis) + return self.library.cumsum( + self.convert_to_typed( + tensor=tensor, + ), + axis, + ) - def conj(self, - tensor + def conj( + self, + tensor, ): """Method to obtain the complex conjugate of a typed tensor. @@ -968,19 +1095,27 @@ def conj(self, Complex conjugate of the typed tensor. """ - return self.library.conj(self.convert_to_typed( - tensor=tensor - )) + return self.library.conj( + self.convert_to_typed( + tensor=tensor, + ), + ) - def min(self, - tensor + def min( + self, + tensor, + axis:int=None, ): - """Method to obtain the minimum value(s) of a typed tensor along an axis. + """Method to obtain the minimum value(s) + of a typed tensor along an axis. Parameters ---------- tensor: Any Given typed tensor. + axis: int, default=None + Given axis. + If ``None`` the overall minimum is returned. Returns ------- @@ -988,12 +1123,21 @@ def min(self, Minimum value(s) of the typed tensor along the given axis. """ - return self.library.min(self.convert_to_typed( - tensor=tensor - )) + return self.library.min( + self.convert_to_typed( + tensor=tensor, + ), + axis, + ) if axis is not None else self.library.min( + self.convert_to_typed( + tensor=tensor, + ), + ) - def max(self, - tensor + def max( + self, + tensor, + axis:int=None, ): """Method to obtain the maximum value(s) of a typed tensor. @@ -1001,6 +1145,9 @@ def max(self, ---------- tensor: Any Given typed tensor. + axis: int, default=None + Given axis. + If ``None`` the overall maximum is returned. Returns ------- @@ -1008,14 +1155,23 @@ def max(self, Maximum value(s) of the typed tensor. """ - return self.library.max(self.convert_to_typed( - tensor=tensor - )) + return self.library.max( + self.convert_to_typed( + tensor=tensor, + ), + axis, + ) if axis is not None else self.library.max( + self.convert_to_typed( + tensor=tensor, + ), + ) - def argmax(self, - tensor + def argmax( + self, + tensor, ): - """Method to obtain the argument of the maximum value(s) of a typed tensor. + """Method to obtain the argument + of the maximum value(s) of a typed tensor. Parameters ---------- @@ -1028,6 +1184,8 @@ def argmax(self, Argument of the maximum value(s) of the typed tensor. """ - return self.library.argmax(self.convert_to_typed( - tensor=tensor - )) + return self.library.argmax( + self.convert_to_typed( + tensor=tensor, + ), + ) diff --git a/quantrl/backends/context_manager.py b/quantrl/backends/context_manager.py index 0f4a217..86ef323 100644 --- a/quantrl/backends/context_manager.py +++ b/quantrl/backends/context_manager.py @@ -6,43 +6,49 @@ __name__ = 'quantrl.backends.context_manager' __authors__ = ["Sampreet Kalita"] __created__ = "2024-10-09" -__updated__ = "2025-05-11" +__updated__ = "2025-05-29" # quantrl modules from .base import BaseBackend -BACKENDS = {} +INSTANCES_BACKEND = {} -def get_backend_instance( +def get_instance_backend( library:str, precision:str='double', - device:str='cuda' - ) -> BaseBackend: + device:str='cuda', +) -> BaseBackend: """Method to obtain an instantiated backend. Parameters ---------- library: str - Name of the library. Options are ``'jax'``, ``'torch'`` and ``'numpy'``. + Name of the library for the backend. + Options are ``'jax'``, ``'torch'`` and ``'numpy'``. precision: str, default='double' - Precision of the numerical values in the backend. Options are ``'single'`` and ``'double'``. + Precision of the numerical values in the backend. + Options are ``'single'`` and ``'double'``. device: str, default='cuda' - Device for the backend. Options are ``'cpu'`` and ``'cuda'``. + Device for the backend. + Options are ``'cpu'`` and ``'cuda'``. Returns ------- Backend: :class:`quantrl.backends.base.BaseBackend` The instantiated backend. """ - if library in BACKENDS: - return BACKENDS[library] + + if library in INSTANCES_BACKEND: + return INSTANCES_BACKEND[library] if 'jax' in library.lower(): try: from .jax import JAXBackend - BACKENDS['jax'] = JAXBackend(precision=precision) + INSTANCES_BACKEND['jax'] = JAXBackend( + precision=precision, + ) library = 'jax' - return BACKENDS[library] + return INSTANCES_BACKEND[library] # use PyTorch if JAX is not installed except ImportError: print("JAX not installed, defaulting to PyTorch") @@ -50,12 +56,18 @@ def get_backend_instance( if 'torch' in library.lower(): from .torch import TorchBackend - BACKENDS['torch'] = TorchBackend(precision=precision, device=device) + INSTANCES_BACKEND['torch'] = TorchBackend( + precision=precision, + device=device, + ) library = 'torch' - return BACKENDS[library] + return INSTANCES_BACKEND[library] - assert 'numpy' in library.lower(), "parameter ``library`` can be either ``'jax'`, ``'torch'`` or ``'numpy'``" + assert 'numpy' in library.lower(), \ + "parameter ``library`` can be either ``'jax'``, ``'torch'`` or ``'numpy'``" from .numpy import NumPyBackend - BACKENDS['numpy'] = NumPyBackend(precision=precision) + INSTANCES_BACKEND['numpy'] = NumPyBackend( + precision=precision, + ) library = 'numpy' - return BACKENDS[library] + return INSTANCES_BACKEND[library] diff --git a/quantrl/backends/jax.py b/quantrl/backends/jax.py index c040ace..c74cdfe 100644 --- a/quantrl/backends/jax.py +++ b/quantrl/backends/jax.py @@ -6,7 +6,7 @@ __name__ = 'quantrl.backends.jax' __authors__ = ["Sampreet Kalita"] __created__ = "2024-03-10" -__updated__ = "2025-05-11" +__updated__ = "2025-06-04" # dependencies from inspect import getfullargspec @@ -24,18 +24,26 @@ class JAXBackend(BaseBackend): """Backend to interface the JAX library. - Refer to :class:`quantrl.backends.base.BaseBackend` for further documentation. + Refer to :class:`quantrl.backends.base.BaseBackend` + for further documentation. + + Parameters + ---------- + precision: str, default='double' + Precision of the numerical values in the backend. + Options are ``'single'`` and ``'double'``. """ - def __init__(self, - precision:str='double' + def __init__( + self, + precision:str='double', ): # initialize BaseBackend super().__init__( name='jax', library=jnp, - tensor_type=jax.Array, - precision=precision + dtype_tensor=jax.Array, + precision=precision, ) # enable 64-bit mode @@ -46,13 +54,13 @@ def __init__(self, self.key = None def transpose( - tensor:jax.Array, - axis_0:int=None, - axis_1:int=None + tensor:jax.Array, + axis_0:int=None, + axis_1:int=None, ) -> jax.Array: if axis_0 is None or axis_1 is None: return self.convert_to_typed( - tensor=tensor + tensor=tensor, ).T # get swapped axes @@ -65,70 +73,76 @@ def transpose( self.jit_transpose = jax.jit( transpose, - static_argnums=(1, 2) + static_argnums=(1, 2), ) self.jit_repeat = jax.jit( lambda tensor, repeats, axis: jnp.repeat(tensor, repeats, axis), - static_argnums=(1, 2) + static_argnums=(1, 2), ) self.jit_add = jax.jit( lambda tensor_0, tensor_1, out: jnp.add(tensor_0, tensor_1), - donate_argnums=(2, ) + donate_argnums=(2, ), ) self.jit_matmul = jax.jit( lambda tensor_0, tensor_1, out: jnp.matmul(tensor_0, tensor_1), - donate_argnums=(2, ) + donate_argnums=(2, ), ) self.jit_dot = jax.jit( lambda tensor_0, tensor_1, out: jnp.dot(tensor_0, tensor_1), - donate_argnums=(2, ) + donate_argnums=(2, ), ) self.jit_concatenate = jax.jit( lambda tensors, axis, out: jnp.concatenate(tensors, axis), static_argnums=(1, ), - donate_argnums=(2, ) + donate_argnums=(2, ), ) self.jit_stack = jax.jit( lambda tensors, axis, out: jnp.stack(tensors, axis), static_argnums=(1, ), - donate_argnums=(2, ) + donate_argnums=(2, ), ) self.jit_update = jax.jit( lambda tensor, indices, values: tensor.at[indices].set(values), - donate_argnums=(0, ) + donate_argnums=(0, ), ) - def convert_to_typed(self, - tensor, - dtype:str=None + def convert_to_typed( + self, + tensor, + dtype:str=None, ) -> jax.Array: if self.is_typed( tensor=tensor, - dtype=dtype + dtype=dtype, ): return tensor - return jnp.array(tensor, dtype=self.dtype_from_str( - dtype=dtype - ) if dtype is not None else None) + return jnp.array( + tensor, + dtype=self.dtype_from_str( + dtype=dtype, + ) if dtype is not None else None, + ) - def convert_to_numpy(self, - tensor, - dtype:str=None + def convert_to_numpy( + self, + tensor, + dtype:str=None, ) -> np.ndarray: return np.asarray(tensor, dtype=self.dtype_from_str( dtype=dtype, - numpy=True + numpy=True, ) if dtype is not None else None) - def generator(self, - seed:int=None + def generator( + self, + seed:int=None, ): if self.key is None: if seed is None: @@ -137,152 +151,202 @@ def generator(self, self.key, key = jax.random.split(self.key) return key - def integers(self, - generator, - shape:tuple, - low:int=0, - high:int=1000, - dtype:str=None + def integers( + self, + generator, + shape:tuple, + low:int=0, + high:int=1000, + dtype:str=None, ) -> jax.Array: - return jnp.asarray(jax.random.randint(generator, shape, low, high), dtype=self.dtype_from_str( - dtype=dtype - )) - - def normal(self, - generator, - shape:tuple, - mean:float=0.0, - std:float=1.0, - dtype:str=None + return jnp.asarray( + jax.random.randint( + generator, + shape, + low, + high, + ), + dtype=self.dtype_from_str( + dtype=dtype, + ), + ) + + def normal( + self, + generator, + shape:tuple, + mean:float=0.0, + std:float=1.0, + dtype:str=None, ) -> jax.Array: - return mean + std * jax.random.normal(generator, shape, dtype=self.dtype_from_str( - dtype=dtype - )) - - def uniform(self, - generator, - shape:tuple, - low:float=0.0, - high:float=1.0, - dtype:str=None + return mean + std * jax.random.normal( + generator, + shape, + dtype=self.dtype_from_str( + dtype=dtype, + ), + ) + + def uniform( + self, + generator, + shape:tuple, + low:float=0.0, + high:float=1.0, + dtype:str=None, ) -> jax.Array: - return jax.random.uniform(generator, shape, minval=low, maxval=high, dtype=self.dtype_from_str( - dtype=dtype - )) - - def transpose(self, - tensor, - axis_0:int=None, - axis_1:int=None + return jax.random.uniform( + generator, + shape, + minval=low, + maxval=high, + dtype=self.dtype_from_str( + dtype=dtype, + ), + ) + + def transpose( + self, + tensor, + axis_0:int=None, + axis_1:int=None, ) -> jax.Array: return self.jit_transpose( tensor=tensor, axis_0=axis_0, - axis_1=axis_1 + axis_1=axis_1, ) - def repeat(self, - tensor, - repeats, - axis + def repeat( + self, + tensor, + repeats, + axis, ) -> jax.Array: return self.jit_repeat( tensor=tensor, repeats=repeats, - axis=axis + axis=axis, ) - def add(self, - tensor_0, - tensor_1, - out + def add( + self, + tensor_0, + tensor_1, + out, ) -> jax.Array: return self.jit_add( tensor_0=tensor_0, tensor_1=tensor_1, - out=out + out=out, ) - def matmul(self, - tensor_0, - tensor_1, - out + def matmul( + self, + tensor_0, + tensor_1, + out, ) -> jax.Array: return self.jit_matmul( tensor_0=tensor_0, tensor_1=tensor_1, - out=out + out=out, ) - def dot(self, - tensor_0, - tensor_1, - out + def dot( + self, + tensor_0, + tensor_1, + out, ) -> jax.Array: return self.jit_dot( tensor_0=tensor_0, tensor_1=tensor_1, - out=out + out=out, ) - def norm(self, - tensor, - axis + def norm( + self, + tensor, + axis, ) -> jax.Array: - return jnp.linalg.norm(tensor, axis=axis) + return jnp.linalg.norm( + tensor, + axis=axis, + ) - def concatenate(self, - tensors:tuple, - axis, - out + def concatenate( + self, + tensors:tuple, + axis, + out, ) -> jax.Array: return self.jit_concatenate( tensors=tensors, axis=axis, - out=out + out=out, ) - def stack(self, - tensors:tuple, - axis, - out + def stack( + self, + tensors:tuple, + axis, + out, ) -> jax.Array: return self.jit_stack( tensors=tensors, axis=axis, - out=out + out=out, ) - def update(self, - tensor, - indices, - values + def update( + self, + tensor, + indices, + values, ) -> jax.Array: return tensor.at[indices].set(values) - def if_else(self, - condition, - func_true, - func_false, - args + def if_else( + self, + condition, + func_true, + func_false, + args, ): - return jax.lax.cond(condition, func_true, func_false, (args, )) + return jax.lax.cond( + condition, + func_true, + func_false, + (args, ), + ) - def iterate_i(self, - func, - iterations_i:int, - Y, - args:tuple=None + def iterate_i( + self, + func, + iterations_i:int, + Y, + args:tuple=None, ) -> jax.Array: # convert to comapatible function _func_args = getfullargspec(func).args - if (_func_args[0] == 'self' and len(_func_args) > 3) or len(_func_args) > 2: + if ( + _func_args[0] == 'self' and len(_func_args) > 3 + ) or len(_func_args) > 2: def body_func(i, state): return (func(i, *state), *state[1:]) else: body_func = func # loop and return typed tensor - return jax.lax.fori_loop(0, iterations_i, body_func, (self.convert_to_typed( - tensor=Y - ), args))[0] + return jax.lax.fori_loop( + 0, + iterations_i, + body_func, + ( + self.convert_to_typed( + tensor=Y, + ), + args, + ), + )[0] diff --git a/quantrl/backends/numpy.py b/quantrl/backends/numpy.py index 18449f3..603aaa0 100644 --- a/quantrl/backends/numpy.py +++ b/quantrl/backends/numpy.py @@ -6,7 +6,7 @@ __name__ = 'quantrl.backends.numpy' __authors__ = ["Sampreet Kalita"] __created__ = "2024-03-10" -__updated__ = "2024-10-13" +__updated__ = "2025-05-29" # dependencies import numpy as np @@ -17,88 +17,124 @@ class NumPyBackend(BaseBackend): """Backend to interface the NumPy library. + Refer to :class:`quantrl.backends.base.BaseBackend` + for further documentation. + Parameters ---------- precision: str, default='double' - Precision of the numerical values in the backend. Options are ``'single'`` and ``'double'``. + Precision of the numerical values in the backend. + Options are ``'single'`` and ``'double'``. """ - def __init__(self, - precision:str='double' + + def __init__( + self, + precision:str='double', ): # initialize BaseBackend super().__init__( name='numpy', library=np, - tensor_type=np.ndarray, + dtype_tensor=np.ndarray, precision=precision ) - def convert_to_typed(self, - tensor, - dtype:str=None + def convert_to_typed( + self, + tensor, + dtype:str=None, ) -> np.ndarray: - return np.asarray(tensor, dtype=self.dtype_from_str( - dtype=dtype, - numpy=True - ) if dtype is not None else None) + return np.asarray( + tensor, + dtype=self.dtype_from_str( + dtype=dtype, + numpy=True, + ) if dtype is not None else None, + ) - def convert_to_numpy(self, - tensor, - dtype:str=None + def convert_to_numpy( + self, + tensor, + dtype:str=None, ) -> np.ndarray: return self.convert_to_typed( tensor=tensor, - dtype=dtype + dtype=dtype, ) - def generator(self, - seed:int=None + def generator( + self, + seed:int=None, ) -> np.random.Generator: if self.seed_sequence is None: self.seed_sequence = self.get_seedsequence(seed) return np.random.default_rng(self.seed_sequence.spawn(1)[0]) - def integers(self, - generator:np.random.Generator, - shape:tuple, - low:int=0, - high:int=1000, - dtype:str=None + def integers( + self, + generator:np.random.Generator, + shape:tuple, + low:int=0, + high:int=1000, + dtype:str=None, ): - return generator.integers(low, high, shape, dtype=self.dtype_from_str( - dtype=dtype - ), endpoint=False) - - def normal(self, - generator:np.random.Generator, - shape:tuple, - mean:float=0.0, - std:float=1.0, - dtype:str=None + return generator.integers( + low, + high, + shape, + dtype=self.dtype_from_str( + dtype=dtype, + ), + endpoint=False, + ) + + def normal( + self, + generator:np.random.Generator, + shape:tuple, + mean:float=0.0, + std:float=1.0, + dtype:str=None, ) -> np.ndarray: - return np.asarray(generator.normal(mean, std, shape), dtype=self.dtype_from_str( - dtype=dtype - )) - - def uniform(self, - generator:np.random.Generator, - shape:tuple, - low:float=0.0, - high:float=1.0, - dtype:str=None + return np.asarray( + generator.normal( + mean, + std, + shape, + ), + dtype=self.dtype_from_str( + dtype=dtype, + ), + ) + + def uniform( + self, + generator:np.random.Generator, + shape:tuple, + low:float=0.0, + high:float=1.0, + dtype:str=None, ) -> np.ndarray: - return np.asarray(generator.uniform(low, high, shape), dtype=self.dtype_from_str( - dtype=dtype - )) - - def transpose(self, - tensor, - axis_0:int=None, - axis_1:int=None + return np.asarray( + generator.uniform( + low, + high, + shape, + ), + dtype=self.dtype_from_str( + dtype=dtype, + ), + ) + + def transpose( + self, + tensor, + axis_0:int=None, + axis_1:int=None, ) -> np.ndarray: if axis_0 is None or axis_1 is None: return self.convert_to_typed( - tensor=tensor + tensor=tensor, ).T # get swapped axes @@ -109,77 +145,114 @@ def transpose(self, return np.transpose(tensor, axes=_axes) - def repeat(self, - tensor, - repeats, - axis + def repeat( + self, + tensor, + repeats, + axis, ) -> np.ndarray: - return np.repeat(tensor, repeats=repeats, axis=axis) + return np.repeat( + tensor, + repeats=repeats, + axis=axis, + ) - def add(self, - tensor_0, - tensor_1, - out + def add( + self, + tensor_0, + tensor_1, + out, ) -> np.ndarray: - return np.add(tensor_0, tensor_1, out=out) + return np.add( + tensor_0, + tensor_1, + out=out, + ) - def matmul(self, - tensor_0, - tensor_1, - out + def matmul( + self, + tensor_0, + tensor_1, + out, ) -> np.ndarray: - return np.matmul(tensor_0, tensor_1, out=out) + return np.matmul( + tensor_0, + tensor_1, + out=out, + ) - def dot(self, - tensor_0, - tensor_1, - out + def dot( + self, + tensor_0, + tensor_1, + out, ) -> np.ndarray: - return np.dot(tensor_0, tensor_1, out=out) + return np.dot( + tensor_0, + tensor_1, + out=out, + ) - def norm(self, - tensor, - axis + def norm( + self, + tensor, + axis, ) -> np.ndarray: - return np.linalg.norm(tensor, axis=axis) + return np.linalg.norm( + tensor, + axis=axis, + ) - def concatenate(self, - tensors:tuple, - axis, - out + def concatenate( + self, + tensors:tuple, + axis, + out, ) -> np.ndarray: - return np.concatenate(tensors, axis=axis, out=out) + return np.concatenate( + tensors, + axis=axis, + out=out, + ) - def stack(self, - tensors:tuple, - axis, - out + def stack( + self, + tensors:tuple, + axis, + out, ) -> np.ndarray: - return np.stack(tensors, axis=axis, out=out) + return np.stack( + tensors, + axis=axis, + out=out, + ) - def update(self, - tensor, - indices, - values + def update( + self, + tensor, + indices, + values, ) -> np.ndarray: tensor[indices] = values return tensor - def if_else(self, - condition, - func_true, - func_false, - args + def if_else( + self, + condition, + func_true, + func_false, + args, ): if condition: return func_true(args) return func_false(args) - def iterate_i(self, - func, - iterations_i:int, - Y, - args:tuple=None + def iterate_i( + self, + func, + iterations_i:int, + Y, + args:tuple=None, ): for i in range(iterations_i): Y = func(i, Y, args) diff --git a/quantrl/backends/torch.py b/quantrl/backends/torch.py index a61fca1..a18abae 100644 --- a/quantrl/backends/torch.py +++ b/quantrl/backends/torch.py @@ -6,7 +6,7 @@ __name__ = 'quantrl.backends.torch' __authors__ = ["Sampreet Kalita"] __created__ = "2024-03-10" -__updated__ = "2025-05-11" +__updated__ = "2025-05-29" # dependencies import numpy as np @@ -18,188 +18,270 @@ class TorchBackend(BaseBackend): """Backend to interface the PyTorch library. + Refer to :class:`quantrl.backends.base.BaseBackend` + for further documentation. + Parameters ---------- precision: str, default='double' - Precision of the numerical values in the backend. Options are ``'single'`` and ``'double'``. + Precision of the numerical values in the backend. + Options are ``'single'`` and ``'double'``. device: str, default='cuda' Device for the backend. Options are ``'cpu'`` and ``'cuda'``. """ - def __init__(self, - precision:str='double', - device:str='cuda' + + def __init__( + self, + precision:str='double', + device:str='cuda', ): # initialize BaseBackend super().__init__( name='torch', library=torch, - tensor_type=torch.Tensor, - precision=precision + dtype_tensor=torch.Tensor, + precision=precision, ) # set default device - assert 'cpu' in device or 'cuda' in device, "Invalid precision opted, options are ``'cpu'`` and ``'cuda'``." + assert 'cpu' in device or 'cuda' in device, \ + "Invalid precision opted, options are ``'cpu'`` and ``'cuda'``." if 'cuda' in device and not torch.cuda.is_available(): print("CUDA not available, defaulting to ``'cpu'``") device = 'cpu' torch.set_default_device(device) self.device = device - def convert_to_typed(self, - tensor, - dtype:str=None + def convert_to_typed( + self, + tensor, + dtype:str=None, ) -> torch.Tensor: if self.is_typed( tensor=tensor, - dtype=dtype + dtype=dtype, ): return tensor - return torch.tensor(tensor, dtype=self.dtype_from_str( - dtype=dtype - ) if dtype is not None else None) + return torch.tensor( + tensor, + dtype=self.dtype_from_str( + dtype=dtype, + ) if dtype is not None else None, + ) - def convert_to_numpy(self, - tensor, - dtype:str=None + def convert_to_numpy( + self, + tensor, + dtype:str=None, ) -> np.ndarray: if self.is_typed( tensor=tensor, - dtype=dtype + dtype=dtype, ): - return np.asarray(tensor.detach().cpu().numpy() if self.device == 'cuda' else tensor.numpy(), dtype=self.dtype_from_str( + return np.asarray( + tensor.detach().cpu().numpy() \ + if self.device == 'cuda' \ + else tensor.numpy(), + dtype=self.dtype_from_str( + dtype=dtype, + numpy=True, + ) if dtype is not None else None, + ) + return np.asarray( + tensor, + dtype=self.dtype_from_str( dtype=dtype, - numpy=True - ) if dtype is not None else None) - return np.asarray(tensor, dtype=self.dtype_from_str( - dtype=dtype, - numpy=True - ) if dtype is not None else None) + numpy=True, + ) if dtype is not None else None, + ) - def generator(self, - seed:int=None + def generator( + self, + seed:int=None, ) -> torch.Generator: if self.seed_sequence is None: self.seed_sequence = self.get_seedsequence(seed) generator = torch.Generator(device=self.device) - generator.manual_seed(int(self.seed_sequence.spawn(1)[0].generate_state(1)[0])) + generator.manual_seed( + int(self.seed_sequence.spawn(1)[0].generate_state(1)[0]), + ) return generator - def integers(self, - generator:torch.Generator, - shape:tuple, - low:int=0, - high:int=1000, - dtype:str=None + def integers( + self, + generator:torch.Generator, + shape:tuple, + low:int=0, + high:int=1000, + dtype:str=None, ) -> torch.Tensor: - return torch.randint(low, high, shape, generator=generator, dtype=self.dtype_from_str( - dtype=dtype - )) - - def normal(self, - generator:torch.Generator, - shape:tuple, - mean:float=0.0, - std:float=1.0, - dtype:str=None + return torch.randint( + low, + high, + shape, + generator=generator, + dtype=self.dtype_from_str( + dtype=dtype, + ), + ) + + def normal( + self, + generator:torch.Generator, + shape:tuple, + mean:float=0.0, + std:float=1.0, + dtype:str=None, ) -> torch.Tensor: return self.empty( shape=shape, - dtype=dtype - ).normal_(mean, std, generator=generator) - - def uniform(self, - generator:torch.Generator, - shape:tuple, - low:float=0.0, - high:float=1.0, - dtype:str=None + dtype=dtype, + ).normal_( + mean, + std, + generator=generator, + ) + + def uniform( + self, + generator:torch.Generator, + shape:tuple, + low:float=0.0, + high:float=1.0, + dtype:str=None, ) -> torch.Tensor: - return low + (high - low) * torch.rand(shape, generator=generator, dtype=self.dtype_from_str( - dtype=dtype - )) - - def transpose(self, - tensor, - axis_0:int=None, - axis_1:int=None + return low + (high - low) * torch.rand( + shape, + generator=generator, + dtype=self.dtype_from_str( + dtype=dtype, + ), + ) + + def transpose( + self, + tensor, + axis_0:int=None, + axis_1:int=None, ) -> torch.Tensor: if axis_0 is None or axis_1 is None: return self.convert_to_typed( - tensor=tensor + tensor=tensor, ).T - return torch.transpose(tensor, dim0=axis_0, dim1=axis_1) + return torch.transpose( + tensor, + dim0=axis_0, + dim1=axis_1, + ) - def repeat(self, - tensor, - repeats, - axis + def repeat( + self, + tensor, + repeats, + axis, ) -> torch.Tensor: - return torch.repeat_interleave(tensor, repeats=repeats, dim=axis) + return torch.repeat_interleave( + tensor, + repeats=repeats, + dim=axis, + ) - def add(self, - tensor_0, - tensor_1, - out + def add( + self, + tensor_0, + tensor_1, + out, ) -> torch.Tensor: - return torch.add(tensor_0, tensor_1, out=out) + return torch.add( + tensor_0, + tensor_1, + out=out, + ) - def matmul(self, - tensor_0, - tensor_1, - out + def matmul( + self, + tensor_0, + tensor_1, + out, ) -> torch.Tensor: - return torch.matmul(tensor_0, tensor_1, out=out) + return torch.matmul( + tensor_0, + tensor_1, + out=out, + ) - def dot(self, - tensor_0, - tensor_1, - out + def dot( + self, + tensor_0, + tensor_1, + out, ) -> torch.Tensor: - return torch.dot(tensor_0, tensor_1, out=out) + return torch.dot( + tensor_0, + tensor_1, + out=out, + ) - def norm(self, - tensor, - axis + def norm( + self, + tensor, + axis, ) -> torch.Tensor: - return torch.norm(tensor, dim=axis) + return torch.norm( + tensor, + dim=axis, + ) - def concatenate(self, - tensors:tuple, - axis, - out + def concatenate( + self, + tensors:tuple, + axis, + out, ) -> torch.Tensor: - return torch.concatenate(tensors, dim=axis, out=out) + return torch.concatenate( + tensors, + dim=axis, + out=out, + ) - def stack(self, - tensors:tuple, - axis, - out + def stack( + self, + tensors:tuple, + axis, + out, ) -> torch.Tensor: - return torch.stack(tensors, dim=axis, out=out) + return torch.stack( + tensors, + dim=axis, + out=out, + ) - def update(self, - tensor, - indices, - values + def update( + self, + tensor, + indices, + values, ) -> torch.Tensor: tensor[indices] = values return tensor - def if_else(self, - condition, - func_true, - func_false, - args + def if_else( + self, + condition, + func_true, + func_false, + args, ): if condition: return func_true(args) return func_false(args) - def iterate_i(self, - func, - iterations_i:int, - Y, - args:tuple=None + def iterate_i( + self, + func, + iterations_i:int, + Y, + args:tuple=None, ): for i in range(iterations_i): Y = func(i, Y, args) diff --git a/quantrl/envs/base.py b/quantrl/envs/base.py index 181c5e0..4e7963e 100644 --- a/quantrl/envs/base.py +++ b/quantrl/envs/base.py @@ -6,7 +6,7 @@ __name__ = 'quantrl.envs.base' __authors__ = ["Sampreet Kalita"] __created__ = "2023-04-25" -__updated__ = "2025-05-11" +__updated__ = "2025-08-20" # dependencies from abc import ABC, abstractmethod @@ -29,10 +29,14 @@ class BaseEnv(ABC): r"""Base environment for reinforcement-learning. - Initializes ``T_norm``, ``T``, ``observation_space``, ``action_space``, ``action_steps``, ``file_path_prefix``, ``io`` and ``plotter``. + Initializes ``T_norm``, ``T``, + ``observation_space``, ``action_space``, ``action_steps``, + ``file_path_prefix``, ``io`` and ``plotter``. - The interfaced environment needs to implement ``_update_states``, ``reset_states``, and ``get_reward`` methods. - Additionally, the ``get_properties`` method should be overridden if ``n_properties`` is non-zero. + The interfaced environment needs to implement ``_update_states``, + ``reset_states``, and ``get_reward`` methods. + Additionally, the ``get_properties`` method should be + overridden if ``n_properties`` is non-zero. Parameters ---------- @@ -55,7 +59,12 @@ class BaseEnv(ABC): action_interval: int Interval at which the actions are updated. Must be positive. data_idxs: list - Indices of the data to store into the ``data`` attribute. The indices can be selected from the complete set of values at each point of time (total ``1 + n_actions + n_observations + n_properties + 1`` elements in the same order, where the first element is the time and the last element is the reward). + Indices of the data to store into the ``data`` attribute. + The indices can be selected from the complete set + of values at each point of time + (total ``1 + n_actions + n_observations + n_properties + 1`` + elements in the same order, where the first element is the time + and the last element is the reward). dir_prefix: str Prefix of the directory where the data will be stored. file_prefix: str @@ -63,29 +72,68 @@ class BaseEnv(ABC): kwargs: dict Keyword arguments. Available options are: - ======================== ================================================ + ======================== ======================================== key value - ======================== ================================================ - has_delay (*bool*) option to implement delay functions. Default is ``False``. - observation_space_range (*list*) range of the observations. Default is ``[-1e12, 1e12]``. - observation_stds (*list* or ``None``) standard deviations of the observed states from the actual states. Default is ``None``. - action_space_range (*list*) range of the actions. The output is scaled by the corresponding action multiplier. Default is ``[-1.0, 1.0]``. - action_space_type (*str*) the type of action space. Options are ``'binary'`` and ``'box'``. Default is ``'box``. - seed (*int*) seed to initialize random number generators. If ``None``, a random integer seed is generated. Default is ``None``. - cache_all_data (*bool*) option to cache all data to disk. If ``False``, only the indices of ``data_idxs`` are stored. Default is ``True``. - cache_dump_interval (*int*) number of environments to cache before dumping to disk. Default is ``100``. - average_over (*int*) number of episodes to run the running average over. This value should be less than or equal to the total number of episodes. Default is ``100``. - plot (*bool*) option to plot the trajectories using ``:class:BaseTrajectoryPlotter``. Default is ``True``. - plot_interval (*int*) number of trajectories after which the plots are updated. Must be a positive integer. Default is ``10``. - plot_idxs (*list*) indices of the data values required to plot at each time step. Default is ``[-1]`` for the cummulative reward. - axes_args (*list*) lists of axis properties. The first element of each is the ``x_label``, the second is ``y_label``, the third is ``[y_limit_min, y_limit_max]`` and the fourth is ``y_scale``. Default is ``[['$t / t_{0}$', '$\\tilde{R}$', [np.sqrt(10) * 1e-1, np.sqrt(10) * 1e6], 'log']]``. - axes_lines_max (*int*) maximum number of lines to display in each plot. Higher numbers slow down the run. Default is ``10``. - axes_cols (*int*) number of columns in the figure. Default is ``2``. - plot_buffer (*bool*) option to store a buffer of plots for to make a gif file. Default is ``False``. - ======================== ================================================ + ======================== ======================================== + has_delay (*bool*) option to implement delay + functions. Default is ``False``. + observation_space_range (*list*) range of the observations. + Default is ``[-1e12, 1e12]``. + observation_stds (*list* or ``None``) standard deviations + of the observed states from + the actual states. Default is ``None``. + action_space_range (*list*) range of the actions. The output + is scaled by the corresponding action + multiplier. Default is ``[-1.0, 1.0]``. + action_space_type (*str*) the type of action space. + Options are ``'binary'`` and ``'box'``. + Default is ``'box'``. + seed (*int*) seed to initialize random number + generators. If ``None``, a random integer + seed is generated. Default is ``None``. + cache_all_data (*bool*) option to cache all data to disk. + If ``False``, only ``data_idxs`` + are stored. Default is ``True``. + cache_dump_interval (*int*) number of environments to cache + before dumping to disk. Default is ``100``. + average_over (*int*) number of episodes to run the + running average over. This value should be + less than or equal to the total number + of episodes. Default is ``100``. + plot (*bool*) option to plot the trajectories + using ``:class:BaseTrajectoryPlotter``. + Default is ``True``. + plot_interval (*int*) number of trajectories after which + the plots are updated. Must be a positive + integer. Default is ``10``. + plot_idxs (*list*) indices of the data values + required to plot at each time step. Default + is ``[-1]`` for the cummulative reward. + axes_args (*list*) lists of axis properties. The + first element of each is the ``x_label``, + the second is ``y_label``, the third is + ``[y_limit_min, y_limit_max]`` and the + fourth is ``y_scale``. Default is + ``[['$t / t_{0}$', '$\\tilde{R}$', + [np.sqrt(10) * 1e-1, np.sqrt(10) * 1e6], + 'log']]``. + axes_lines_max (*int*) maximum number of lines to display + in each plot. Higher numbers slow down + the run. Default is ``10``. + axes_cols (*int*) number of columns in the figure. + Default is ``2``. + plot_buffer (*bool*) option to store a buffer of plots + for to make a gif file. + Default is ``False``. + ======================== ======================================== """ - default_axis_args_learning_curve = ['Episodes', 'Average Return', [np.sqrt(10) * 1e-4, np.sqrt(10) * 1e6], 'log'] + default_axis_args_learning_curve = [ + 'Episodes', + 'Average Return', + [np.sqrt(10) * 1e-4, np.sqrt(10) * 1e6], + 'log', + ] """list: Default axis arguments to plot the learning curve.""" base_env_kwargs = { @@ -102,28 +150,34 @@ class BaseEnv(ABC): 'plot_interval': 10, 'plot_idxs': [-1], 'axes_args': [ - ['$t / \\tau$', '$\\tilde{R}$', [np.sqrt(10) * 1e-5, np.sqrt(10) * 1e4], 'log'] + [ + '$t / \\tau$', + '$\\tilde{R}$', + [np.sqrt(10) * 1e-5, np.sqrt(10) * 1e4], + 'log', + ], ], 'axes_lines_max': 10, 'axes_cols': 2, - 'plot_buffer': False + 'plot_buffer': False, } """dict: Default values of all keyword arguments.""" - def __init__(self, - backend:BaseBackend, - t_norm_max:float, - t_norm_ssz:float, - t_norm_mul:float, - n_observations:int, - n_properties:int, - n_actions:int, - action_maximums:list, - action_interval:int, - data_idxs:list, - dir_prefix:str, - file_prefix:str, - **kwargs + def __init__( + self, + backend:BaseBackend, + t_norm_max:float, + t_norm_ssz:float, + t_norm_mul:float, + n_observations:int, + n_properties:int, + n_actions:int, + action_maximums:list, + action_interval:int, + data_idxs:list, + dir_prefix:str, + file_prefix:str, + **kwargs, ): """Class constructor for BaseEnv.""" @@ -132,37 +186,59 @@ def __init__(self, kwargs[key] = kwargs.get(key, self.base_env_kwargs[key]) # validate arguments - assert t_norm_max > t_norm_ssz, "maximum normalized time should be greater than the normalized step size" - assert n_properties >= 0, "parameter ``n_properties`` should be non-negative" - assert action_interval > 0, "parameter ``action_interval`` should be a positive integer" - assert len(data_idxs) > 0, "parameter ``data_idxs`` should be a list containing at least one element" - assert kwargs['observation_stds'] is None or isinstance(kwargs['observation_stds'], list), "parameter ``observation_stds`` should be a list" - assert kwargs['seed'] is None or isinstance(kwargs['seed'], int), "parameter ``seed`` should be an integer or ``None``" - assert isinstance(kwargs['cache_all_data'], bool), "parameter ``cache_all_data`` should be a boolean" - assert kwargs['plot_interval'] > 0, "parameter ``plot_interval`` should be a positive integer" - assert len(kwargs['plot_idxs']) == len(kwargs['axes_args']), "number of indices for plot should match number of axes arguments" - assert len(kwargs['observation_space_range']) == 2, "parameter ``observation_space_range`` should contain two elements for the minimum and maximum values, both inclusive" - assert len(kwargs['action_space_range']) == 2, "parameter ``action_space_range`` should contain two elements for the minimum and maximum values, both inclusive" - assert kwargs['action_space_type'] in ['binary', 'box'], "parameter ``action_space_type`` can be either ``'binary'`` or ``'box'``" - assert kwargs['cache_dump_interval'] > 0, "parameter ``cache_dump_interval`` should be a positive integer" + assert t_norm_max > t_norm_ssz, \ + "maximum normalized time should be greater than the normalized step size" + assert n_properties >= 0, \ + "parameter ``n_properties`` should be non-negative" + assert action_interval > 0, \ + "parameter ``action_interval`` should be a positive integer" + assert len(data_idxs) > 0, \ + "parameter ``data_idxs`` should be a list containing at least one element" + assert kwargs['observation_stds'] is None \ + or isinstance(kwargs['observation_stds'], list), \ + "parameter ``observation_stds`` should be a list" + assert kwargs['seed'] is None or isinstance(kwargs['seed'], int), \ + "parameter ``seed`` should be an integer or ``None``" + assert isinstance(kwargs['cache_all_data'], bool), \ + "parameter ``cache_all_data`` should be a boolean" + assert kwargs['plot_interval'] > 0, \ + "parameter ``plot_interval`` should be a positive integer" + assert len(kwargs['plot_idxs']) == len(kwargs['axes_args']), \ + "number of indices for plot should match number of axes arguments" + assert len(kwargs['observation_space_range']) == 2, \ + "parameter ``observation_space_range`` should contain " \ + + "two elements for the minimum and maximum values, " \ + + "both inclusive" + assert len(kwargs['action_space_range']) == 2, \ + "parameter ``action_space_range`` should contain " \ + + "two elements for the minimum and maximum values, " \ + + "both inclusive" + assert kwargs['action_space_type'] in ['binary', 'box'], \ + "parameter ``action_space_type`` can be either ``'binary'`` or ``'box'``" + assert kwargs['cache_dump_interval'] > 0, \ + "parameter ``cache_dump_interval`` should be a positive integer" # set backend self.backend = backend # frequently used variables - self.numpy_int = self.backend.dtypes['numpy'][self.backend.precision]['integer'] - self.numpy_real = self.backend.dtypes['numpy'][self.backend.precision]['real'] + self.numpy_int = self.backend.dtypes['numpy']\ + [self.backend.precision]['integer'] + self.numpy_real = self.backend.dtypes['numpy']\ + [self.backend.precision]['real'] # time attributes self.t_norm_max = t_norm_max self.t_norm_ssz = t_norm_ssz self.t_norm_mul = t_norm_mul # truncate before maximum time if not divisible - self.shape_T = (self.numpy_int(self.t_norm_max / self.t_norm_ssz) + 1, ) - self.T_norm = np.arange(self.shape_T[0], dtype=self.numpy_real) * self.t_norm_ssz + self.shape_T = \ + (self.numpy_int(self.t_norm_max / self.t_norm_ssz) + 1, ) + self.T_norm = np.arange(self.shape_T[0], dtype=self.numpy_real) \ + * self.t_norm_ssz self.T = self.backend.convert_to_typed( tensor=self.T_norm, - dtype='real' + dtype='real', ) * t_norm_mul self.t_ssz = self.t_norm_ssz * t_norm_mul @@ -173,13 +249,13 @@ def __init__(self, low=self.observation_space_range[0], high=self.observation_space_range[1], shape=(self.n_observations, ), - dtype=self.numpy_real + dtype=self.numpy_real, ) self.observation_stds = kwargs['observation_stds'] if self.observation_stds is not None: self.observation_stds = self.backend.convert_to_typed( tensor=self.observation_stds, - dtype='real' + dtype='real', ) # property attributes @@ -201,14 +277,16 @@ def __init__(self, low=self.action_space_range[0], high=self.action_space_range[1], shape=(self.n_actions, ), - dtype=self.numpy_real + dtype=self.numpy_real, ) self.action_maximums = self.backend.convert_to_typed( tensor=action_maximums, dtype='integer' if self.action_space_type == 'binary' else 'real' ) self.action_interval = action_interval - self.action_steps = self.numpy_int(np.ceil((self.shape_T[0] - 1) / self.action_interval)) + self.action_steps = self.numpy_int( + np.ceil((self.shape_T[0] - 1) / self.action_interval) + ) # align delay with action interval self.has_delay = kwargs['has_delay'] @@ -218,15 +296,19 @@ def __init__(self, self.seed = kwargs['seed'] # data constants - self.dir_path = dir_prefix + '/' + '_'.join([ + self.dir_path = dir_prefix + "/" + "_".join([ str(t_norm_max), str(t_norm_ssz), str(t_norm_mul), str(action_maximums), - str(action_interval) + str(action_interval), ]) - self.file_path_prefix = self.dir_path + '/' + file_prefix - self.n_data = 1 + self.n_actions + self.n_observations + self.n_properties + 1 + self.file_path_prefix = self.dir_path + "/" + file_prefix + self.n_data = 1 \ + + self.n_actions \ + + self.n_observations \ + + self.n_properties \ + + 1 self.average_over = self.numpy_int(kwargs['average_over']) # initialize IO @@ -235,7 +317,7 @@ def __init__(self, self.cache_dump_interval = kwargs['cache_dump_interval'] self.io = FileIO( disk_cache_dir=self.file_path_prefix + '_cache', - cache_dump_interval=self.cache_dump_interval + cache_dump_interval=self.cache_dump_interval, ) # plot constants @@ -250,7 +332,7 @@ def __init__(self, axes_lines_max=kwargs['axes_lines_max'], axes_cols=kwargs['axes_cols'], show_title=True, - save_dir=self.file_path_prefix + '_plots' + save_dir=self.file_path_prefix + '_plots', ) # initialize buffers @@ -275,7 +357,9 @@ def _update_states(self): Returns ------- States: Any - The updated states with shape either ``(action_interval + 1, n_observations)`` or ``(action_interval + 1, n_envs, n_observations). + The updated states with shape either + ``(action_interval + 1, n_observations)`` or + ``(action_interval + 1, n_envs, n_observations). """ raise NotImplementedError @@ -287,7 +371,11 @@ def reset_states(self): Returns ------- states_0: Any - The initial states with shape either ``(n_observations, )`` or ``(n_envs, n_observations)``, which are assigned to all elements of ``Observations`` with shape ``(action_interval + 1, n_observations, )`` or ``(action_interval + 1, n_envs, n_observations)``. + The initial states with shape either + ``(n_observations, )`` or ``(n_envs, n_observations)``, + which are assigned to all elements of ``Observations`` + with shape ``(action_interval + 1, n_observations, )`` or + ``(action_interval + 1, n_envs, n_observations)``. """ raise NotImplementedError @@ -298,7 +386,9 @@ def get_properties(self): Returns ------- Properties: Any - The properties calculated from ``Observations`` with shape either ``(action_interval + 1, n_properties)`` or ``(action_interval + 1, n_envs, n_properties)``. + The properties calculated from ``Observations`` + with shape either ``(action_interval + 1, n_properties)`` or + ``(action_interval + 1, n_envs, n_properties)``. """ raise NotImplementedError @@ -310,15 +400,18 @@ def get_reward(self): Returns ------- Reward: Any - The reward calculated using ``States``, ``Observations`` or ``Properties`` with shape either ``(action_interval + 1, )`` or ``(action_interval + 1, n_envs)``. + The reward calculated using ``States``, + ``Observations`` or ``Properties`` with shape either + ``(action_interval + 1, )`` or ``(action_interval + 1, n_envs)``. """ raise NotImplementedError - def validate_base(self, - shape_reset_states:tuple, - shape_get_properties:tuple, - shape_get_reward:tuple + def validate_base( + self, + shape_reset_states:tuple, + shape_get_properties:tuple, + shape_get_reward:tuple, ): """Method to validate the base environment. @@ -335,25 +428,26 @@ def validate_base(self, try: # validate initial states states_0 = self.backend.convert_to_typed( - tensor=self.reset_states() + tensor=self.reset_states(), ) assert self.backend.shape( tensor=states_0 - ) == shape_reset_states, f"``reset_states`` should return an array with shape ``{shape_reset_states}``" + ) == shape_reset_states, \ + f"``reset_states`` should return an array with shape ``{shape_reset_states}``" # initialize states self.States = self.backend.repeat( tensor=self.backend.reshape( tensor=states_0, - shape=(1, *shape_reset_states) + shape=(1, *shape_reset_states), ), repeats=self.action_interval + 1, - axis=0 + axis=0, ) # initialize observations self.Observations = self.backend.update( tensor=self.Observations, indices=(slice(None), ), - values=self.States + values=self.States, ) # validate properties if self.n_properties > 0: @@ -362,20 +456,27 @@ def validate_base(self, ) assert self.backend.shape( tensor=self.Properties - ) == shape_get_properties, f"``get_properties`` should return an array with shape ``{shape_get_properties}``" + ) == shape_get_properties, \ + "``get_properties`` should return an array " \ + + f"with shape ``{shape_get_properties}``" # validate reward self.Reward = self.backend.convert_to_typed( tensor=self.get_reward() ) assert self.backend.shape( tensor=self.Reward - ) == shape_get_reward, f"``get_reward`` should return an array with shape ``{shape_get_reward}``" + ) == shape_get_reward, \ + f"``get_reward`` should return an array with shape ``{shape_get_reward}``" except AttributeError as error: - print(f"Missing required method or attribute: ({error}). Refer to **Notes** of :class:`quantrl.envs.base.BaseEnv` for the implementation format of the missing method or add the missing attribute to the ``reset_states`` method.") + print(f"Missing required method or attribute: ({error}) " \ + + "Refer to **Notes** of :class:`quantrl.envs.base.BaseEnv` " \ + + "for the implementation format of the missing method or " \ + + "add the missing attribute to the ``reset_states`` method.") sys.exit() def reset(self): - """Method to reset the time and obtain initial states as a typed tensor. + """Method to reset the time and obtain initial states + as a typed tensor. Returns ------- @@ -390,24 +491,24 @@ def reset(self): # initialize states states_0 = self.backend.convert_to_typed( - tensor=self.reset_states() + tensor=self.reset_states(), ) _shape = self.backend.shape( - tensor=states_0 + tensor=states_0, ) self.States = self.backend.repeat( tensor=self.backend.reshape( tensor=states_0, - shape=(1, *_shape) + shape=(1, *_shape), ), repeats=self.action_interval + 1, - axis=0 + axis=0, ) # initialize measurement noises if self.observation_stds is not None: self.Observation_noises = self.backend.normal( generator=self.backend.generator( - seed=self.seed + seed=self.seed, ), shape=(self.shape_T[0], *_shape), mean=0.0, @@ -416,28 +517,32 @@ def reset(self): ) * self.backend.repeat( tensor=self.backend.reshape( tensor=self.observation_stds, - shape=(1, *_shape) + shape=(1, *_shape), ), repeats=self.shape_T[0], - axis=0 + axis=0, ) # initialize observations - observations_0 = states_0 + (self.Observation_noises[0] if self.observation_stds is not None else 0.0) + observations_0 = states_0 + (self.Observation_noises[0] \ + if self.observation_stds is not None \ + else 0.0) self.Observations = self.backend.repeat( tensor=self.backend.reshape( tensor=observations_0, shape=(1, *self.backend.shape( - tensor=observations_0 + tensor=observations_0, )) ), repeats=self.action_interval + 1, - axis=0 + axis=0, ) return observations_0 def update(self): - """Method to update the time, observations, properties and reward and obtain the final set of observations and reward as typed tensors. + """Method to update the time, observations, + properties and reward and obtain the + final set of observations and reward as typed tensors. Returns ------- @@ -451,24 +556,31 @@ def update(self): # set evaluation times _dim_T = self.shape_T[0] - self.T_step = self.T[self.t_idx:self.t_idx + self.action_interval + 1] if self.t_idx + self.action_interval < _dim_T else self.T[self.t_idx:] + self.T_step = \ + self.T[self.t_idx:self.t_idx + self.action_interval + 1] \ + if self.t_idx + self.action_interval < _dim_T \ + else self.T[self.t_idx:] _dim_T_step = self.backend.shape( - tensor=self.T_step + tensor=self.T_step, )[0] # update actual states and observed states self.States = self._update_states() - self.Observations = self.States + (self.Observation_noises[self.t_idx:self.t_idx + _dim_T_step] if self.observation_stds is not None else 0.0) + self.Observations = self.States + ( + self.Observation_noises[self.t_idx:self.t_idx + _dim_T_step] \ + if self.observation_stds is not None \ + else 0.0 + ) # update properties if self.n_properties > 0: self.Properties = self.backend.convert_to_typed( - tensor=self.get_properties() + tensor=self.get_properties(), ) # update rewards self.Reward = self.backend.convert_to_typed( - tensor=self.get_reward() + tensor=self.get_reward(), ) # update time @@ -479,7 +591,11 @@ def update(self): # check if completed terminated = not self.t_idx + 1 < _dim_T - return self.Observations[_dim_T_step - 1], self.Reward[_dim_T_step - 1], terminated + return ( + self.Observations[_dim_T_step - 1], + self.Reward[_dim_T_step - 1], + terminated, + ) def check_truncation(self): """Method to check if the current episode needs to be truncated. @@ -492,43 +608,52 @@ def check_truncation(self): # check if out of bounds return bool(self.backend.max( - tensor=self.Observations + tensor=self.Observations, ) > self.observation_space_range[1] or self.backend.min( - tensor=self.Observations + tensor=self.Observations, ) < self.observation_space_range[0]) - def plot_learning_curve(self, - data_rewards:np.ndarray=None, - n_episodes:int=None, - axis_args:list=None, - hold:bool=False + def plot_learning_curve( + self, + data_rewards:np.ndarray=None, + n_episodes:int=None, + axis_args:list=None, + hold:bool=False, ): """Method to plot the learning curve. - Either one of the parameters ``n_episodes`` or ``data_rewards`` should be provided. + Either one of the parameters ``n_episodes`` + or ``data_rewards`` should be provided. Parameters ---------- data_rewards: :class:`numpy.ndarray`, default=None - Cummulative rewards with shape ``(n_trajectories, 1)``. Loads data from disk cache if ``None``. + Cummulative rewards with shape ``(n_trajectories, 1)``. + Loads data from disk cache if ``None``. n_episodes: int, default=None Total number of episodes to load from cache. axis_args: list, default=None - Axis properties. The first element is the ``x_label``, the second is ``y_label``, the third is ``[y_limit_min, y_limit_max]`` and the fourth is ``y_scale``. + Axis properties. The first element is the ``x_label``, + the second is ``y_label``, + the third is ``[y_limit_min, y_limit_max]`` and + the fourth is ``y_scale``. hold: bool, default=False Option to hold the plot. """ # validate arguments - assert data_rewards is not None or n_episodes is not None, "either one of the parameters ``data_rewards`` or ``n_episodes`` should be provided" + assert data_rewards is not None or n_episodes is not None, \ + "either one of the parameters ``data_rewards`` or ``n_episodes`` should be provided" # extract frequently used variables _idx_s = self._idx_s - _idx_e = self._idx_s + data_rewards.shape[0] - 1 if data_rewards is not None else n_episodes - 1 - file_name = self.file_path_prefix + '_' + '_'.join([ - 'learning_curve', + _idx_e = self._idx_s + data_rewards.shape[0] - 1 \ + if data_rewards is not None \ + else n_episodes - 1 + file_name = self.file_path_prefix + "_" + "_".join([ + "learning_curve", str(_idx_s), - str(_idx_e) + str(_idx_e), ]) # get reward data from file @@ -546,21 +671,25 @@ def plot_learning_curve(self, # initialize plotter plotter = LearningCurvePlotter( - axis_args=axis_args if axis_args is not None and len(axis_args) == 4 else self.default_axis_args_learning_curve, - average_over=self.average_over if self.average_over < data_rewards.shape[0] else int(data_rewards.shape[0] / 2) + axis_args=axis_args \ + if axis_args is not None and len(axis_args) == 4 \ + else self.default_axis_args_learning_curve, + average_over=self.average_over \ + if self.average_over < data_rewards.shape[0] \ + else int(data_rewards.shape[0] / 2), ) # update plot plotter.add_data( data_rewards=data_rewards, - renew=False + renew=False, ) # save plot self.io.save_data( data=data_rewards, - file_name=file_name + file_name=file_name, ) plotter.save_plot( - file_name=file_name + file_name=file_name, ) # hold plot if hold: @@ -569,12 +698,13 @@ def plot_learning_curve(self, # close plotter plotter.close() - def replay_trajectories(self, - n_episodes, - idx_start:int=0, - plot_interval:int=0, - make_gif:bool=True, - hold:bool=False + def replay_trajectories( + self, + n_episodes, + idx_start:int=0, + plot_interval:int=0, + make_gif:bool=True, + hold:bool=False, ): """Method to replay trajectories in a given range. @@ -585,20 +715,24 @@ def replay_trajectories(self, idx_start: int, default=0 Starting index for the cached files. plot_interval: int, default=0 - Number of trajectories after which the plots are updated. If non-positive, the environment's ``plot_interval`` value is taken. + Number of trajectories after which the plots are updated. + If non-positive, the environment's + ``plot_interval`` value is taken. make_gif: bool, default=True Option to create a gif file for the replay. """ # extract frequently used variables _idx_e = n_episodes - 1 - _interval = self.plot_interval if plot_interval <= 0 else plot_interval + _interval = self.plot_interval \ + if plot_interval <= 0 \ + else plot_interval # get replay data in the given range replay_data = self.io.get_disk_cache( idx_start=idx_start, idx_end=_idx_e, - idxs=self.plot_idxs + idxs=self.plot_idxs, ) # update plotter @@ -607,22 +741,22 @@ def replay_trajectories(self, desc="Plotting", leave=False, mininterval=0.5, - disable=False + disable=False, ): self.plotter.plot_lines( xs=self.T_norm, Y=replay_data[i], traj_idx=idx_start + i, - update_buffer=True + update_buffer=True, ) # make gif if make_gif: self.plotter.make_gif( - file_name=self.file_path_prefix + '_' + '_'.join([ - 'replay', + file_name=self.file_path_prefix + "_" + "_".join([ + "replay", str(idx_start), str(_idx_e), - str(_interval) + str(_interval), ]) ) # hold plot @@ -632,9 +766,10 @@ def replay_trajectories(self, # close plotter self.plotter.close() - def close_base(self, - n_episodes, - save_replay=True + def close_base( + self, + n_episodes, + save_replay=True, ): """Method to close the base environment. @@ -649,18 +784,19 @@ def close_base(self, if self.plot and save_replay: # make replay gif self.plotter.make_gif( - file_name=self.file_path_prefix + '_' + '_'.join([ - 'replay', + file_name=self.file_path_prefix + "_" + "_".join([ + "replay", str(self._idx_s), str(n_episodes - 1), - str(self.plot_interval) + str(self.plot_interval), ]) ) # close plotter self.plotter.close() # clean - del self.T, self.T_norm, self.T_step, self.States, self.Observations, self.Reward + del self.T, self.T_norm, self.T_step, \ + self.States, self.Observations, self.Reward if self.n_properties > 0: del self.Properties del self @@ -671,20 +807,21 @@ class BaseGymEnv(BaseEnv, Env): Refer to :class:`quantrl.envs.base.BaseEnv` for its documentation. """ - def __init__(self, - backend:BaseBackend, - t_norm_max:float, - t_norm_ssz:float, - t_norm_mul:float, - n_observations:int, - n_properties:int, - n_actions:int, - action_maximums:list, - action_interval:int, - data_idxs:list, - dir_prefix:str, - file_prefix:str, - **kwargs + def __init__( + self, + backend:BaseBackend, + t_norm_max:float, + t_norm_ssz:float, + t_norm_mul:float, + n_observations:int, + n_properties:int, + n_actions:int, + action_maximums:list, + action_interval:int, + data_idxs:list, + dir_prefix:str, + file_prefix:str, + **kwargs, ): """Class constructor for BaseGymEnv.""" @@ -702,7 +839,7 @@ def __init__(self, data_idxs=data_idxs, dir_prefix=dir_prefix, file_prefix=file_prefix, - **kwargs + **kwargs, ) # initialize Gymnasium environment @@ -713,19 +850,19 @@ def __init__(self, self.actions = None self.States = self.backend.empty( shape=(self.action_interval + 1, self.n_observations), - dtype='real' + dtype='real', ) self.Observations = self.backend.empty( shape=(self.action_interval + 1, self.n_observations), - dtype='real' + dtype='real', ) self.Properties = self.backend.empty( shape=(self.action_interval + 1, self.n_properties), - dtype='real' + dtype='real', ) self.Reward = self.backend.empty( shape=(self.action_interval + 1, ), - dtype='real' + dtype='real', ) self.rewards = None self.data_rewards = [] @@ -737,15 +874,20 @@ def validate(self): return super().validate_base( shape_reset_states=(self.n_observations, ), - shape_get_properties=(self.action_interval + 1, self.n_properties), - shape_get_reward=(self.action_interval + 1, ) + shape_get_properties=( + self.action_interval + 1, + self.n_properties, + ), + shape_get_reward=(self.action_interval + 1, ), ) - def reset(self, - seed:float=None, - options:dict=None + def reset( + self, + seed:float=None, + options:dict=None, ): - """Method to reset all variables for a new trajectory and obtain the initial observations as a NumPy array or a typed tensor. + """Method to reset all variables for a new trajectory and + obtain the initial observations as a NumPy array or a typed tensor. Parameters ---------- @@ -765,10 +907,16 @@ def reset(self, # update buffers self.traj_idx += 1 self.rewards = 0.0 - self.all_data = np.zeros((self.shape_T[0], self.n_data), dtype=self.numpy_real) + self.all_data = np.zeros(( + self.shape_T[0], + self.n_data, + ), dtype=self.numpy_real) # store selected data - self.data = np.zeros((self.shape_T[0], len(self.data_idxs)), dtype=self.numpy_real) + self.data = np.zeros(( + self.shape_T[0], + len(self.data_idxs), + ), dtype=self.numpy_real) # reset variables observations_0 = super().reset() @@ -777,10 +925,12 @@ def reset(self, 'traj_idx': self.traj_idx } - def step(self, - action + def step( + self, + action, ): - """Method to take one single step and obtain the observations and reward as NumPy arrays or typed tensors. + """Method to take one single step and obtain the + observations and reward as NumPy arrays or typed tensors. Parameters ---------- @@ -804,7 +954,7 @@ def step(self, # set actions self.actions = self.backend.convert_to_typed( tensor=action, - dtype='real' + dtype='real', ) * self.action_maximums # get observations, properties and reward @@ -816,13 +966,13 @@ def step(self, # check if truncation required truncated = self.check_truncation() if truncated > 0: - print(f'Trajectory #{self.traj_idx} truncated') + print(f"Trajectory #{self.traj_idx} truncated") # if trajectory ends if terminated or truncated: # update cache self.io.update_cache( - data=self.all_data if self.cache_all_data else self.data + data=self.all_data if self.cache_all_data else self.data, ) # update episode reward self.data_rewards.append(self.rewards) @@ -832,7 +982,7 @@ def step(self, xs=self.T_norm, Y=self.all_data[:, self.plot_idxs], traj_idx=self.traj_idx, - update_buffer=self.plot_buffer + update_buffer=self.plot_buffer, ) return observations, reward, terminated, truncated, {} @@ -849,7 +999,7 @@ def update_data(self): # frequently used variables _dim = self.backend.shape( - tensor=self.T_step + tensor=self.T_step, )[0] # update rewards @@ -860,33 +1010,44 @@ def update_data(self): # extract values _Ts_step = self.backend.reshape( tensor=self.T_step, - shape=(_dim, 1) + shape=(_dim, 1), ) _actions = self.backend.repeat( tensor=self.backend.reshape( tensor=self.actions, - shape=(1, self.n_actions) + shape=(1, self.n_actions), ), repeats=_dim, - axis=0 + axis=0, ) _Rewards = self.backend.repeat( tensor=self.backend.convert_to_typed( tensor=[[self.rewards]], - dtype='real' + dtype='real', ), repeats=_dim, - axis=0 + axis=0, ) # concatenate values - _tensors = (_Ts_step, _actions, self.Observations, self.Properties, _Rewards) if self.n_properties > 0 else (_Ts_step, _actions, self.Observations, _Rewards) + _tensors = ( + _Ts_step, + _actions, + self.Observations, + self.Properties, + _Rewards, + ) if self.n_properties > 0 else ( + _Ts_step, + _actions, + self.Observations, + _Rewards, + ) _data_backend = self.backend.concatenate( tensors=_tensors, axis=1, - out=None + out=None, ) _data = self.backend.convert_to_numpy( - tensor=_data_backend + tensor=_data_backend, ) # update data @@ -896,9 +1057,10 @@ def update_data(self): # clear cache del _data_backend, _tensors, _Ts_step, _actions, _Rewards - def evolve(self, - show_progress=True, - close=True + def evolve( + self, + show_progress=True, + close=True, ): """Method to freely evolve the trajectory. @@ -919,7 +1081,7 @@ def evolve(self, desc="Evolving", leave=True, mininterval=0.5, - disable=not show_progress + disable=not show_progress, ): # set actions self.actions = self.action_maximums @@ -938,7 +1100,7 @@ def evolve(self, # udpate cache self.io.update_cache( - data=self.all_data if self.cache_all_data else self.data + data=self.all_data if self.cache_all_data else self.data, ) # update episode reward self.data_rewards.append(self.rewards) @@ -948,41 +1110,47 @@ def evolve(self, xs=self.T_norm, Y=self.all_data[:, self.plot_idxs], traj_idx=self.traj_idx, - update_buffer=self.plot_buffer + update_buffer=self.plot_buffer, ) # close environment if close: self.reset() self.close( - save=False + save=False, ) - def close(self, - save=True, - axis_args=None + def close( + self, + save=True, + axis_args=None, ): """Method to close the environment. Parameters ---------- save: bool, default=True - Option to save the learning curve and make a gif file from the plot buffer. + Option to save the learning curve and + make a gif file from the plot buffer. axis_args: list, default=None - Axis properties for the learning curve. The first element is the ``x_label``, the second is ``y_label``, the third is ``[y_limit_min, y_limit_max]`` and the fourth is ``y_scale``. + Axis properties for the learning curve. + The first element is the ``x_label``, + the second is ``y_label``, + the third is ``[y_limit_min, y_limit_max]`` and + the fourth is ``y_scale``. """ if save: _data_rewards = self.backend.convert_to_numpy( tensor=self.data_rewards, - dtype='real' + dtype='real', ) # save learning curve self.plot_learning_curve( data_rewards=_data_rewards.ravel(), n_episodes=None, axis_args=axis_args, - hold=False + hold=False, ) del self.rewards, self.data_rewards, self.all_data, self.data @@ -998,29 +1166,32 @@ def close(self, ) class BaseSB3Env(BaseEnv, VecEnv): - r"""Stable-Baselines3-based vectorized environments for reinforcement-learning. + r"""Stable-Baselines3-based vectorized environments + for reinforcement-learning. Initializes ``action_maximums_batch``. Refer to :class:`quantrl.envs.base.BaseEnv` for its documentation. - The additional parameter ``n_envs`` denotes the number of environments to run in parallel and overrides the ``cache_dump_interval`` parameter. + The additional parameter ``n_envs`` denotes the number of environments + to run in parallel and overrides the ``cache_dump_interval`` parameter. """ - def __init__(self, - backend:BaseBackend, - t_norm_max:float, - t_norm_ssz:float, - t_norm_mul:float, - n_envs:int, - n_observations:int, - n_properties:int, - n_actions:int, - action_maximums:list, - action_interval:int, - data_idxs:list, - dir_prefix:str, - file_prefix:str, - **kwargs + def __init__( + self, + backend:BaseBackend, + t_norm_max:float, + t_norm_ssz:float, + t_norm_mul:float, + n_envs:int, + n_observations:int, + n_properties:int, + n_actions:int, + action_maximums:list, + action_interval:int, + data_idxs:list, + dir_prefix:str, + file_prefix:str, + **kwargs, ): """Class constructor for BaseSB3Env.""" @@ -1039,7 +1210,7 @@ def __init__(self, dir_prefix=dir_prefix, file_prefix=file_prefix, cache_dump_interval=n_envs, - **kwargs + **kwargs, ) # update attributes @@ -1048,17 +1219,17 @@ def __init__(self, self.action_maximums_batch = self.backend.repeat( tensor=self.backend.reshape( tensor=self.action_maximums, - shape=(1, self.n_actions) + shape=(1, self.n_actions), ), repeats=self.n_envs, - axis=0 + axis=0, ) # initialize SB3 environment VecEnv.__init__(self, num_envs=self.n_envs, observation_space=self.observation_space, - action_space=self.action_space + action_space=self.action_space, ) # initialize buffers @@ -1066,19 +1237,19 @@ def __init__(self, self.actions = None self.States = self.backend.empty( shape=(self.action_interval + 1, self.n_envs, self.n_observations), - dtype='real' + dtype='real', ) self.Observations = self.backend.empty( shape=(self.action_interval + 1, self.n_envs, self.n_observations), - dtype='real' + dtype='real', ) self.Properties = self.backend.empty( shape=(self.action_interval + 1, self.n_envs, self.n_properties), - dtype='real' + dtype='real', ) self.Reward = self.backend.empty( shape=(self.action_interval + 1, self.n_envs), - dtype='real' + dtype='real', ) self.rewards = None self.data_rewards = [] @@ -1090,23 +1261,36 @@ def validate(self): """Method to validate BaseSB3Env.""" return super().validate_base( - shape_reset_states=(self.n_envs, self.n_observations), - shape_get_properties=(self.action_interval + 1, self.n_envs, self.n_properties), - shape_get_reward=(self.action_interval + 1, self.n_envs) + shape_reset_states=( + self.n_envs, + self.n_observations, + ), + shape_get_properties=( + self.action_interval + 1, + self.n_envs, + self.n_properties, + ), + shape_get_reward=( + self.action_interval + 1, + self.n_envs, + ), ) - def env_is_wrapped(self, - wrapper_class, - indices=None + def env_is_wrapped( + self, + wrapper_class, + indices=None, ): - """Method to check if a batch of sub-environments are wrapped with the given wrapper. + """Method to check if a batch of sub-environments + are wrapped with the given wrapper. Parameters ---------- wrapper_class: :class:`gymnasium.Wrapper` Wrapper class. indices: int or list, default=None - Indices of the environments. If ``None``, the values for all sub-environments are considered. + Indices of the environments. + If ``None``, the values for all sub-environments are considered. Returns ------- @@ -1114,13 +1298,15 @@ def env_is_wrapped(self, Whether the batch of sub-environments are wrapped. """ - return [env_util.is_wrapped(self, wrapper_class) for _ in range(indices if indices is not None else self.n_envs)] + return [env_util.is_wrapped(self, wrapper_class) \ + for _ in range(indices if indices is not None else self.n_envs)] - def env_method(self, - method_name, - *method_args, - indices=None, - **method_kwargs + def env_method( + self, + method_name, + *method_args, + indices=None, + **method_kwargs, ): """Method to call other methods of the sub-environments. @@ -1131,7 +1317,8 @@ def env_method(self, method_args: tuple Additional positional arguments. indices: int or list, default=None - Indices of the environments. If ``None``, the values for all sub-environments are considered. + Indices of the environments. + If ``None``, the values for all sub-environments are considered. method_kwargs: dict Additional keyword arguments. @@ -1141,11 +1328,13 @@ def env_method(self, Methods of the sub-environments. """ - return [getattr(self, method_name)(*method_args, **method_kwargs) for _ in range(indices if indices is not None else self.n_envs)] + return [getattr(self, method_name)(*method_args, **method_kwargs) \ + for _ in range(indices if indices is not None else self.n_envs)] - def get_attr(self, - attr_name, - indices=None + def get_attr( + self, + attr_name, + indices=None, ): """Method to obtain attributes of the sub-environments. @@ -1154,7 +1343,8 @@ def get_attr(self, attr_name: str Name of the attribute. indices: int or list, default=None - Indices of the environments. If ``None``, the values for all sub-environments are considered. + Indices of the environments. + If ``None``, the values for all sub-environments are considered. Returns ------- @@ -1162,12 +1352,14 @@ def get_attr(self, Attributes of the sub-environments. """ - return [getattr(self, attr_name) for _ in range(indices if indices is not None else self.n_envs)] + return [getattr(self, attr_name) \ + for _ in range(indices if indices is not None else self.n_envs)] - def set_attr(self, - attr_name, - value, - indices=None + def set_attr( + self, + attr_name, + value, + indices=None, ): """Method to assign attributes of the sub-environments. @@ -1178,16 +1370,20 @@ def set_attr(self, value: any Value of the attribute. indices: int or list, default=None - Indices of the environments. If ``None``, the values for all sub-environments are considered. + Indices of the environments. + If ``None``, the values for all sub-environments are considered. """ - return [setattr(self, attr_name, value) for _ in range(indices if indices is not None else self.n_envs)] + return [setattr(self, attr_name, value) \ + for _ in range(indices if indices is not None else self.n_envs)] - def reset(self, - seed:float=None, - options:dict=None + def reset( + self, + seed:float=None, + options:dict=None, ): - """Method to reset all variables for a new batch and obtain the initial observations as a NumPy array or a typed tensor. + """Method to reset all variables for a new batch and + obtain the initial observations as a NumPy array or a typed tensor. Parameters ---------- @@ -1208,7 +1404,8 @@ def reset(self, return observations_0 def _reset(self): - """Method to reset all variables for a new batch and obtain the initial observations as a typed tensor. + """Method to reset all variables for a new batch and + obtain the initial observations as a typed tensor. Returns ------- @@ -1220,19 +1417,31 @@ def _reset(self): self.batch_idx += 1 self.rewards = self.backend.zeros( shape=(self.n_envs, ), - dtype='real' + dtype='real', ) # store selected data and plot data - self.data = np.zeros((self.n_envs, self.shape_T[0], len(self.data_idxs)), dtype=self.numpy_real) + self.data = np.zeros(( + self.n_envs, + self.shape_T[0], + len(self.data_idxs), + ), dtype=self.numpy_real) if self.plot: - self.plotter_env_idxs = self.env_idx_arr[(self.batch_idx * self.n_envs + self.env_idx_arr) % self.plot_interval == 0] - self.plotter_env_data = np.zeros((len(self.plotter_env_idxs), self.shape_T[0], len(self.plot_idxs)), dtype=self.numpy_real) + self.plotter_env_idxs = self.env_idx_arr[ + (self.batch_idx * self.n_envs + self.env_idx_arr) \ + % self.plot_interval == 0 + ] + self.plotter_env_data = np.zeros(( + len(self.plotter_env_idxs), + self.shape_T[0], + len(self.plot_idxs), + ), dtype=self.numpy_real) return super().reset() - def step_async(self, - actions + def step_async( + self, + actions, ): """Method to prepare for one single step. @@ -1272,7 +1481,7 @@ def step_wait(self): # check if truncation required truncated = self.check_truncation() if truncated > 0: - print(f'Batch #{self.batch_idx} truncated') + print(f"Batch #{self.batch_idx} truncated") # if trajectory ends if terminated or truncated: @@ -1282,8 +1491,9 @@ def step_wait(self): self.plotter.plot_lines( xs=self.T_norm, Y=self.plotter_env_data[i, :, :], - traj_idx=self.batch_idx * self.n_envs + plotter_env_idx, - update_buffer=self.plot_buffer + traj_idx=self.batch_idx * self.n_envs \ + + plotter_env_idx, + update_buffer=self.plot_buffer, ) # update episode reward self.data_rewards.append(self.rewards) @@ -1291,7 +1501,8 @@ def step_wait(self): # reset variables observations = self._reset() - return observations, reward, [terminated or truncated] * self.n_envs, [{}] * self.n_envs + return observations, reward, \ + [terminated or truncated] * self.n_envs, [{}] * self.n_envs def update_data(self): """Method to update the batch data for the step. @@ -1305,7 +1516,7 @@ def update_data(self): # frequently used variables _dim = self.backend.shape( - tensor=self.T_step + tensor=self.T_step, )[0] # update rewards @@ -1318,45 +1529,56 @@ def update_data(self): shape=(1, _dim, 1) ), repeats=self.n_envs, - axis=0 + axis=0, ) _Observations = self.backend.transpose( tensor=self.Observations[:_dim], axis_0=1, - axis_1=0 + axis_1=0, ) _Properties = None if self.n_properties > 0: _Properties = self.backend.transpose( tensor=self.Properties[:_dim], axis_0=1, - axis_1=0 + axis_1=0, ) _actions = self.backend.repeat( tensor=self.backend.reshape( tensor=self.actions, - shape=(self.n_envs, 1, self.n_actions) + shape=(self.n_envs, 1, self.n_actions), ), repeats=_dim, - axis=1 + axis=1, ) _Rewards = self.backend.repeat( tensor=self.backend.reshape( tensor=self.rewards, - shape=(self.n_envs, 1, 1) + shape=(self.n_envs, 1, 1), ), repeats=_dim, - axis=1 + axis=1, ) # concatenate values - _tensors = (_Ts_step, _actions, _Observations, _Properties, _Rewards) if self.n_properties > 0 else (_Ts_step, _actions, _Observations, _Rewards) + _tensors = ( + _Ts_step, + _actions, + _Observations, + _Properties, + _Rewards, + ) if self.n_properties > 0 else ( + _Ts_step, + _actions, + _Observations, + _Rewards, + ) _data_backend = self.backend.concatenate( tensors=_tensors, axis=2, - out=None + out=None, ) _data = self.backend.convert_to_numpy( - tensor=_data_backend + tensor=_data_backend, ) # dump part data to disk @@ -1364,7 +1586,7 @@ def update_data(self): self.io.dump_part_async( data=_data, batch_idx=self.batch_idx, - part_idx=self.action_idx - 1 + part_idx=self.action_idx - 1, ) # update selected data @@ -1373,17 +1595,20 @@ def update_data(self): self.data[:, _idx_start:_idx_stop, :] = _data[:, :, self.data_idxs] # update plot data if self.plot: - self.plotter_env_data[:, _idx_start:_idx_stop, :] = _data[self.plotter_env_idxs][:, :, self.plot_idxs] + self.plotter_env_data[:, _idx_start:_idx_stop, :] = \ + _data[self.plotter_env_idxs][:, :, self.plot_idxs] # clear cache - del _data_backend, _tensors, _Ts_step, _actions, _Observations, _Rewards + del _data_backend, _tensors, _Ts_step, \ + _actions, _Observations, _Rewards if self.n_properties > 0: del _Properties - def evolve(self, - show_progress=True, - close=True, - save=False + def evolve( + self, + show_progress=True, + close=True, + save=False, ): """Method to freely evolve the trajectory. @@ -1406,7 +1631,7 @@ def evolve(self, desc="Evolving", leave=True, mininterval=0.5, - disable=not show_progress + disable=not show_progress, ): # set actions self.actions = self.action_maximums_batch @@ -1432,13 +1657,14 @@ def evolve(self, desc="Plotting", leave=True, mininterval=0.5, - disable=False + disable=False, ): self.plotter.plot_lines( xs=self.T_norm, Y=self.plotter_env_data[_i, :, :], - traj_idx=self.batch_idx * self.n_envs + self.plotter_env_idxs[_i], - update_buffer=self.plot_buffer + traj_idx=self.batch_idx * self.n_envs \ + + self.plotter_env_idxs[_i], + update_buffer=self.plot_buffer, ) self.plotter.hold_plot() @@ -1449,31 +1675,37 @@ def evolve(self, save=save ) - def close(self, - save=True, - axis_args=None + def close( + self, + save=True, + axis_args=None, ): """Method to close the environment. Parameters ---------- save: bool, default=True - Option to save the learning curve and make a gif file from the plot buffer. + Option to save the learning curve and + make a gif file from the plot buffer. axis_args: list, default=None - Axis properties for the learning curve. The first element is the ``x_label``, the second is ``y_label``, the third is ``[y_limit_min, y_limit_max]`` and the fourth is ``y_scale``. + Axis properties for the learning curve. + The first element is the ``x_label``, + the second is ``y_label``, + the third is ``[y_limit_min, y_limit_max]`` and + the fourth is ``y_scale``. """ # save learning curve if save: _data_rewards = self.backend.convert_to_numpy( tensor=self.data_rewards, - dtype='real' + dtype='real', ) self.plot_learning_curve( data_rewards=_data_rewards.ravel(), n_episodes=None, axis_args=axis_args, - hold=False + hold=False, ) # clear cache @@ -1483,11 +1715,11 @@ def close(self, # close io self.io.close( - dump_cache=True + dump_cache=True, ) # clean super().close_base( n_episodes=self.batch_idx, - save_replay=save + save_replay=save, ) diff --git a/quantrl/envs/deterministic.py b/quantrl/envs/deterministic.py index 0801fca..60cb479 100644 --- a/quantrl/envs/deterministic.py +++ b/quantrl/envs/deterministic.py @@ -6,25 +6,36 @@ __name__ = 'quantrl.envs.deterministic' __authors__ = ["Sampreet Kalita"] __created__ = "2023-04-25" -__updated__ = "2025-05-11" +__updated__ = "2025-08-20" # quantrl modules -from ..backends.context_manager import get_backend_instance -from ..solvers.context_manager import get_IVP_solver +from ..backends.context_manager import get_instance_backend +from ..solvers.context_manager import get_solver_ivp from .base import BaseGymEnv, BaseSB3Env # TODO: ABC for common processes class LinearizedHOEnv(BaseGymEnv): - """Class to interface deterministic linearized harmonic oscillator environments. - - Initializes ``dim_corrs``, ``num_corrs``, ``A``, ``D``, ``is_A_constant``, ``is_D_constant`` and ``solver``. - The interfaced environment requires ``default_params`` dictionary defined before initializing the parent class. - - The interfaced environment needs to implement ``reset_states`` and ``get_reward`` methods. - Additionally, the ``get_properties`` method should be overridden if ``n_properties`` is non-zero. - Refer to **Notes** of :class:`quantrl.envs.base.BaseEnv` for their implementations. - The default ``func`` method can be used to call ``get_mode_rates`` for rates of change of the classical mode amplitudes, ``get_A`` for the Jacobian of the quantum fluctuation quadratures and ``get_D`` for the quantum noise correlations by overriding the corresponding methods. + """Class to interface deterministic linearized + harmonic oscillator environments. + + Initializes ``dim_corrs``, ``num_corrs``, ``A``, ``D``, + ``is_A_constant``, ``is_D_constant`` and ``solver``. + The interfaced environment requires ``default_params`` + dictionary defined before initializing the parent class. + + The interfaced environment needs to implement + ``reset_states`` and ``get_reward`` methods. + Additionally, the ``get_properties`` method + should be overridden if ``n_properties`` is non-zero. + Refer to **Notes** of :class:`quantrl.envs.base.BaseEnv` + for their implementations. + + The default ``func`` method can be used to call + ``get_mode_rates`` for rates of change of the classical mode amplitudes, + ``get_A`` for the Jacobian of the quantum fluctuation quadratures and + ``get_D`` for the quantum noise correlations + by overriding the corresponding methods. Parameters ---------- @@ -53,26 +64,49 @@ class LinearizedHOEnv(BaseGymEnv): action_maximums: list Maximum values of each action. action_interval: int - Interval at which the actions are updated. Must be positive. + Interval at which the actions are updated. + Must be positive. data_idxs: list - Indices of the data to store into the ``data`` attribute. The indices can be selected from the complete set of values at each point of time (total ``1 + n_actions + n_observations + n_properties + 1`` elements in the same order, where the first element is the time and the last element is the reward). + Indices of the data to store into the ``data`` attribute. + The indices can be selected from the + complete set of values at each point of time + (total ``1 + n_actions + n_observations + n_properties + 1`` + elements in the same order, where the first element + is the time and the last element is the reward). backend_library: str, default='numpy' - Solver to use for each step. Options are ``'jax'`` for JAX-based solvers, ``'torch'`` for PyTorch-based solvers and ``'numpy'`` for NumPy/SciPy-based solvers. + Solver to use for each step. + Options are ``'jax'`` for JAX-based solvers, + ``'torch'`` for PyTorch-based solvers and + ``'numpy'`` for NumPy/SciPy-based solvers. backend_precision: str, default='double' - Precision of the numerical values in the backend. Options are ``'single'`` and ``'double'``. + Precision of the numerical values in the backend. + Options are ``'single'`` and ``'double'``. backend_device: str, default='cuda' Device to run the solver. Options are ``'cpu'`` and ``'cuda'``. - dir_prefix: str, default='data' + dir_prefix: str, default="data" Prefix of the directory where the data will be stored. kwargs: dict, optional - Keyword arguments. Refer to the ``kwargs`` parameter of :class:`quantrl.envs.base.BaseEnv` for available options. Additional options are: + Keyword arguments. Refer to the ``kwargs`` parameter of + :class:`quantrl.envs.base.BaseEnv` for available options. + Additional options are: ============ ================================================ key value ============ ================================================ - ode_method (*str*) method used to solve the ODEs/DDEs. Available options are ``'dopri5'``, ``'dopri8'`` and ``'tsit5'`` for a Diffrax-based solver, ``'adaptive_huen'``, ``'bosh3'``, ``'dopri5'``, ``'dopri8'``, ``'fehlberg2'`` and ``'tsit5'`` for a TorchDiffEq-based solver and ``'BDF'``, ``'DOP853'``, ``'LSODA'``, ``'Radau'``, ``'RK23'``, ``'RK45'``, ``'dop853'``, ``'dopri5'``, ``'lsoda'``, ``'vode'`` and ``'zvode'`` for a SciPy-based solver. Default is ``'dopri5'``. - ode_atol (*float*) absolute tolerance of the ODE/DDE solver. Default is ``1e-9``. - ode_rtol (*float*) relative tolerance of the ODE/DDE solver. Default is ``1e-6``. + ode_method (*str*) method used to solve the ODEs/DDEs. + Available options are ``'dopri5'``, ``'dopri8'`` + and ``'tsit5'`` for a Diffrax-based solver, + ``'adaptive_huen'``, ``'bosh3'``, ``'dopri5'``, + ``'dopri8'``, ``'fehlberg2'`` and ``'tsit5'`` + for a TorchDiffEq-based solver and + ``'BDF'``, ``'DOP853'``, ``'LSODA'``, ``'Radau'``, + ``'RK23'``, ``'RK45'``, ``'dop853'``, ``'dopri5'``, + ``'lsoda'``, ``'vode'`` and ``'zvode'`` + for a SciPy-based solver. Default is ``'dopri5'``. + ode_atol (*float*) absolute tolerance of the ODE/DDE solver. + Default is ``1e-9``. + ode_rtol (*float*) relative tolerance of the ODE/DDE solver. + Default is ``1e-6``. ============ ================================================ """ @@ -82,47 +116,49 @@ class LinearizedHOEnv(BaseGymEnv): default_ode_solver_params = { 'ode_method': 'dopri5', 'ode_atol': 1e-9, - 'ode_rtol': 1e-6 + 'ode_rtol': 1e-6, } """dict: Default parameters of the ODE solver.""" backend_libraries = ['jax', 'torch', 'numpy'] """list: Available backend libraries.""" - def __init__(self, - name:str, - desc:str, - params:dict, - num_modes:int, - num_quads:int, - t_norm_max:float, - t_norm_ssz:float, - t_norm_mul:float, - n_observations:int, - n_properties:int, - n_actions:int, - action_maximums:list, - action_interval:int, - data_idxs:list, - backend_library:str='numpy', - backend_precision:str='double', - backend_device:str='cuda', - dir_prefix:str='data', - **kwargs + def __init__( + self, + name:str, + desc:str, + params:dict, + num_modes:int, + num_quads:int, + t_norm_max:float, + t_norm_ssz:float, + t_norm_mul:float, + n_observations:int, + n_properties:int, + n_actions:int, + action_maximums:list, + action_interval:int, + data_idxs:list, + backend_library:str='numpy', + backend_precision:str='double', + backend_device:str='cuda', + dir_prefix:str="data", + **kwargs, ): """Class constructor for LinearizedHOEnv.""" # validate arguments - assert backend_library in self.backend_libraries, f"parameter ``backend_library`` should be one of ``{self.backend_libraries}``" + assert backend_library in self.backend_libraries, \ + f"parameter ``backend_library`` should be one of ``{self.backend_libraries}``" # select backend - backend = get_backend_instance( + backend = get_instance_backend( library=backend_library, precision=backend_precision, - device=backend_device + device=backend_device, ) - IVPSolver = get_IVP_solver( - library=backend_library + IVPSolver = get_solver_ivp( + library=backend_library, ) # set constants @@ -142,11 +178,11 @@ def __init__(self, # set matrices self.A = backend.zeros( shape=self.dim_corrs, - dtype='real' + dtype='real', ) self.D = backend.zeros( shape=self.dim_corrs, - dtype='real' + dtype='real', ) self.is_A_constant = False self.is_D_constant = False @@ -163,9 +199,11 @@ def __init__(self, action_maximums=action_maximums, action_interval=action_interval, data_idxs=data_idxs, - dir_prefix=(dir_prefix if dir_prefix != 'data' else ('data/' + self.name.lower()) + '/env'), - file_prefix='lho_env', - **kwargs + dir_prefix=(dir_prefix \ + if dir_prefix != "data" \ + else ("data/" + self.name.lower()) + "/env"), + file_prefix="lho_env", + **kwargs, ) # initialize solver @@ -178,80 +216,92 @@ def __init__(self, 'atol': kwargs['ode_atol'], 'rtol': kwargs['ode_rtol'], 'is_stiff': False, - 'step_interval': self.action_interval + 'step_interval': self.action_interval, }, func_controls=getattr(self, 'func_controls', None), has_delay=self.has_delay, func_delay=getattr(self, 'func_delay', None), delay_interval=self.action_interval, - backend=self.backend + backend=self.backend, ) # initialize buffers if self.num_modes != 0: self.mode_rates_real = self.backend.zeros( shape=(2 * self.num_modes, ), - dtype='real' + dtype='real', ) if self.num_corrs != 0: self.matmul_0 = self.backend.empty( shape=self.dim_corrs, - dtype='real' + dtype='real', ) self.matmul_1 = self.backend.empty( shape=self.dim_corrs, - dtype='real' + dtype='real', ) self.sum_0 = self.backend.empty( shape=self.dim_corrs, - dtype='real' + dtype='real', ) self.sum_1 = self.backend.empty( shape=self.dim_corrs, - dtype='real' + dtype='real', ) self.y_rates = self.backend.empty( shape=(2 * self.num_modes + self.num_corrs, ), - dtype='real' + dtype='real', ) def _update_states(self): return self.solver.step( T_step=self.T_step, y_0=self.States[-1], - params=self.actions + params=self.actions, ) - def func(self, - t, - y, - args:tuple + def func( + self, + t, + y, + args:tuple, ): - """Method to obtain the rates of change of the real-valued modes and correlations. + """Method to obtain the rates of change + of the real-valued modes and correlations. Parameters ---------- t: float Time at which the values are calculated. y: Any - Real-valued modes and flattened correlations with shape ``(2 * num_modes + num_corrs, )``. First ``num_modes`` elements contain the real parts of the modes, the next ``num_modes`` elements contain the imaginary parts of the modes, and the last ``num_corrs`` elements contain the correlations. When ``num_modes`` is ``0``, only the correlations are included. When ``num_corrs`` is ``0``, only the modes are included. + Real-valued modes and flattened correlations + with shape ``(2 * num_modes + num_corrs, )``. + First ``num_modes`` elements contain the real parts of the modes, + the next ``num_modes`` elements contain the imaginary parts + of the modes, and the last ``num_corrs`` elements + contain the correlations. + When ``num_modes`` is ``0``, only the correlations are included. + When ``num_corrs`` is ``0``, only the modes are included. args: tuple Actions, control function and delay function. Returns ------- rates: Any - Rates of change of the real-valued modes and flattened correlations with shape ``(2 * num_modes + num_corrs, )``. + Rates of change of the real-valued modes + and flattened correlations with shape + ``(2 * num_modes + num_corrs, )``. """ # extract frequently used variables if self.num_modes != 0: - modes = y[:self.num_modes] + 1.0j * y[self.num_modes:2 * self.num_modes] + modes = y[:self.num_modes] \ + + 1.0j * y[self.num_modes:2 * self.num_modes] # get real-valued mode rates _mode_rates_real = self.get_mode_rates_real( t=t, modes_real=y[:2 * self.num_modes], - args=args + args=args, ) if self.num_corrs == 0: return _mode_rates_real @@ -261,21 +311,21 @@ def func(self, if self.num_corrs != 0: corrs = self.backend.reshape( tensor=y[2 * self.num_modes:], - shape=self.dim_corrs + shape=self.dim_corrs, ) # get drift matrix A = self.A if self.is_A_constant else self.get_A( t=t, modes=modes, - args=args + args=args, ) # get noise matrix D = self.D if self.is_D_constant else self.get_D( t=t, modes=modes, - args=args + args=args, ) # get flattened correlation rates @@ -285,21 +335,21 @@ def func(self, tensor_0=self.backend.matmul( tensor_0=A, tensor_1=corrs, - out=self.matmul_0 + out=self.matmul_0, ), tensor_1=self.backend.matmul( tensor_0=corrs, tensor_1=self.backend.transpose( tensor=A, axis_0=0, - axis_1=1 + axis_1=1, ), - out=self.matmul_1 + out=self.matmul_1, ), - out=self.sum_0 + out=self.sum_0, ), tensor_1=D, - out=self.sum_1 + out=self.sum_1, ) ) @@ -309,16 +359,17 @@ def func(self, return self.backend.concatenate( tensors=( _mode_rates_real, - _corr_rates_flat + _corr_rates_flat, ), axis=0, - out=self.y_rates + out=self.y_rates, ) - def get_A(self, - t, - modes, - args:tuple + def get_A( + self, + t, + modes, + args:tuple, ): """Method to obtain the Jacobian of quantum fluctuation quadratures. @@ -334,15 +385,17 @@ def get_A(self, Returns ------- A: Any - Jacobian of quantum fluctuation quadratures with shape ``(num_quads, num_quads)``. + Jacobian of quantum fluctuation quadratures + with shape ``(num_quads, num_quads)``. """ raise NotImplementedError - def get_D(self, - t, - modes, - args:tuple + def get_D( + self, + t, + modes, + args:tuple, ): """Method to obtain the quantum noise correaltions. @@ -363,12 +416,14 @@ def get_D(self, raise NotImplementedError - def get_mode_rates(self, - t, - modes, - args + def get_mode_rates( + self, + t, + modes, + args, ): - """Method to obtain the rates of change of the classical mode amplitudes. + """Method to obtain the rates of change + of the classical mode amplitudes. Parameters ---------- @@ -382,33 +437,39 @@ def get_mode_rates(self, Returns ------- D: Any - Rates of change of the classical mode amplitudes with shape ``(num_modes, )``. + Rates of change of the classical mode amplitudes + with shape ``(num_modes, )``. """ raise NotImplementedError - def get_mode_rates_real(self, - t, - modes_real, - args:tuple + def get_mode_rates_real( + self, + t, + modes_real, + args:tuple, ): - """Method to obtain the real-valued rates of change of the classical mode amplitudes. + """Method to obtain the real-valued rates of change + of the classical mode amplitudes. - The interfaced environment needs to implement the ``get_mode_rates`` method. + The interfaced environment needs to + implement the ``get_mode_rates`` method. Parameters ---------- t: float Time at which the values are calculated. modes_real: Any - Real-valued classical mode amplitudes with shape ``(2 * num_modes, )``. + Real-valued classical mode amplitudes + with shape ``(2 * num_modes, )``. args: tuple Actions, control function and delay function. Returns ------- mode_rates_real: Any - Real-valued rates of change of the classcial mode amplitudes with shape ``(2 * num_modes, )``. + Real-valued rates of change of the classcial mode + amplitudes with shape ``(2 * num_modes, )``. """ # handle null @@ -418,37 +479,52 @@ def get_mode_rates_real(self, # get complex-valued mode rates mode_rates = self.get_mode_rates( t=t, - modes=modes_real[:self.num_modes] + 1.0j * modes_real[self.num_modes:], - args=args + modes=modes_real[:self.num_modes] \ + + 1.0j * modes_real[self.num_modes:], + args=args, ) # return real-valued mode rates return self.backend.concatenate( tensors=( self.backend.real( - tensor=mode_rates + tensor=mode_rates, ), self.backend.imag( - tensor=mode_rates + tensor=mode_rates, ) ), axis=0, - out=self.mode_rates_real + out=self.mode_rates_real, ) class LinearizedHOVecEnv(BaseSB3Env): - """Class to interface deterministic linearized harmonic oscillator vectorized environments. + """Class to interface deterministic linearized + harmonic oscillator vectorized environments. - Initializes ``dim_corrs``, ``num_corrs``, ``A``, ``D``, ``is_A_constant``, ``is_D_constant`` and ``solver``. - The interfaced environment requires ``default_params`` dictionary defined before initializing the parent class. + Initializes ``dim_corrs``, ``num_corrs``, ``A``, ``D``, + ``is_A_constant``, ``is_D_constant`` and ``solver``. + The interfaced environment requires ``default_params`` + dictionary defined before initializing the parent class. The paramter ``n_envs`` overrides the ``cache_dump_interval`` parameter. - For massively parallel models, the ``data_idxs`` parameter should be carefully selected to initialize the ``data`` attribute with shape ``(n_envs, t_dim, n_data_idxs)``. - In such cases, it is advisable to use the JAX backend with the ``cache_all_data`` paramter to ``False``. - - The interfaced environment needs to implement ``reset_states`` and ``get_reward`` methods. - Additionally, the ``get_properties`` method should be overridden if ``n_properties`` is non-zero. - Refer to **Notes** of :class:`quantrl.envs.base.BaseEnv` for their implementations. - The default ``func`` method can be used to call ``get_mode_rates`` for rates of change of the classical mode amplitudes, ``get_A`` for the Jacobian of the quantum fluctuation quadratures and ``get_D`` for the quantum noise correlations by overriding the corresponding methods. + For massively parallel models, the ``data_idxs`` parameter should be + carefully selected to initialize the ``data`` attribute + with shape ``(n_envs, t_dim, n_data_idxs)``. + In such cases, it is advisable to use the JAX backend + with the ``cache_all_data`` paramter to ``False``. + + The interfaced environment needs to implement + ``reset_states`` and ``get_reward`` methods. + Additionally, the ``get_properties`` method + should be overridden if ``n_properties`` is non-zero. + Refer to **Notes** of :class:`quantrl.envs.base.BaseEnv` + for their implementations. + + The default ``func`` method can be used to call + ``get_mode_rates`` for rates of change of the classical mode amplitudes, + ``get_A`` for the Jacobian of the quantum fluctuation quadratures and + ``get_D`` for the quantum noise correlations + by overriding the corresponding methods. Parameters ---------- @@ -469,7 +545,8 @@ class LinearizedHOVecEnv(BaseSB3Env): t_norm_mul: float Multiplier to revert the normalization. n_envs: int - Number of environments to run in parallel. This value overrides the optional ``cache_dump_interval`` parameter. + Number of environments to run in parallel. + This value overrides the optional ``cache_dump_interval`` parameter. n_observations: int Total number of observations. n_properties: int @@ -479,26 +556,49 @@ class LinearizedHOVecEnv(BaseSB3Env): action_maximums: list Maximum values of each action. action_interval: int - Interval at which the actions are updated. Must be positive. + Interval at which the actions are updated. + Must be positive. data_idxs: list - Indices of the data to store into the ``data`` attribute. The indices can be selected from the complete set of values at each point of time (total ``1 + n_actions + n_observations + n_properties + 1`` elements in the same order, where the first element is the time and the last element is the reward). + Indices of the data to store into the ``data`` attribute. + The indices can be selected from the + complete set of valuesat each point of time + (total ``1 + n_actions + n_observations + n_properties + 1`` + elements in the same order, where the first element + is the time and the last element is the reward). backend_library: str, default='numpy' - Solver to use for each step. Options are ``'jax'`` for JAX-based solvers, ``'torch'`` for PyTorch-based solvers and ``'numpy'`` for NumPy/SciPy-based solvers. + Solver to use for each step. + Options are ``'jax'`` for JAX-based solvers, + ``'torch'`` for PyTorch-based solvers and + ``'numpy'`` for NumPy/SciPy-based solvers. backend_precision: str, default='double' - Precision of the numerical values in the backend. Options are ``'single'`` and ``'double'``. + Precision of the numerical values in the backend. + Options are ``'single'`` and ``'double'``. backend_device: str, default='cuda' Device to run the solver. Options are ``'cpu'`` and ``'cuda'``. - dir_prefix: str, default='data' + dir_prefix: str, default="data" Prefix of the directory where the data will be stored. kwargs: dict - Keyword arguments. Refer to the ``kwargs`` parameter of :class:`quantrl.envs.base.BaseEnv` for available options. Additional options are: + Keyword arguments. Refer to the ``kwargs`` parameter of + :class:`quantrl.envs.base.BaseEnv` for available options. + Additional options are: ============ ================================================ key value ============ ================================================ - ode_method (*str*) method used to solve the ODEs/DDEs. Available options are ``'dopri8'``, ``'dopri5'`` and ``'tsit5'`` for a Diffrax-based solver, ``'dopri8'``, ``'dopri5'``, ``'bosh3'``, ``'fehlberg2'`` and ``'adaptive_huen'`` for a TorchDiffEq-based solver and ``'BDF'``, ``'DOP853'``, ``'LSODA'``, ``'Radau'``, ``'RK23'``, ``'RK45'``, ``'dop853'``, ``'dopri5'``, ``'lsoda'``, ``'zvode'`` and ``'vode'`` for a SciPy-based solver. Default is ``'vode'``. - ode_atol (*float*) absolute tolerance of the ODE/DDE solver. Default is ``1e-9``. - ode_rtol (*float*) relative tolerance of the ODE/DDE solver. Default is ``1e-6``. + ode_method (*str*) method used to solve the ODEs/DDEs. + Available options are ``'dopri5'``, ``'dopri8'`` + and ``'tsit5'`` for a Diffrax-based solver, + ``'adaptive_huen'``, ``'bosh3'``, ``'dopri5'``, + ``'dopri8'``, ``'fehlberg2'`` and ``'tsit5'`` + for a TorchDiffEq-based solver and + ``'BDF'``, ``'DOP853'``, ``'LSODA'``, ``'Radau'``, + ``'RK23'``, ``'RK45'``, ``'dop853'``, ``'dopri5'``, + ``'lsoda'``, ``'vode'`` and ``'zvode'`` + for a SciPy-based solver. Default is ``'dopri5'``. + ode_atol (*float*) absolute tolerance of the ODE/DDE solver. + Default is ``1e-9``. + ode_rtol (*float*) relative tolerance of the ODE/DDE solver. + Default is ``1e-6``. ============ ================================================ """ @@ -508,48 +608,50 @@ class LinearizedHOVecEnv(BaseSB3Env): default_ode_solver_params = { 'ode_method': 'dopri5', 'ode_atol': 1e-9, - 'ode_rtol': 1e-6 + 'ode_rtol': 1e-6, } """dict: Default parameters of the ODE solver.""" backend_libraries = ['jax', 'torch', 'numpy'] """list: Available backend libraries.""" - def __init__(self, - name:str, - desc:str, - params:dict, - num_modes:int, - num_quads:int, - t_norm_max:float, - t_norm_ssz:float, - t_norm_mul:float, - n_envs:int, - n_observations:int, - n_properties:int, - n_actions:int, - action_maximums:list, - action_interval:int, - data_idxs:list, - backend_library:str='numpy', - backend_precision:str='double', - backend_device:str='cuda', - dir_prefix:str='data', - **kwargs + def __init__( + self, + name:str, + desc:str, + params:dict, + num_modes:int, + num_quads:int, + t_norm_max:float, + t_norm_ssz:float, + t_norm_mul:float, + n_envs:int, + n_observations:int, + n_properties:int, + n_actions:int, + action_maximums:list, + action_interval:int, + data_idxs:list, + backend_library:str='numpy', + backend_precision:str='double', + backend_device:str='cuda', + dir_prefix:str="data", + **kwargs, ): """Class constructor for LinearizedHOEnv.""" # validate arguments - assert backend_library in self.backend_libraries, f"parameter ``backend_library`` should be one of ``{self.backend_libraries}``" + assert backend_library in self.backend_libraries, \ + f"parameter ``backend_library`` should be one of ``{self.backend_libraries}``" # select backend - backend = get_backend_instance( + backend = get_instance_backend( library=backend_library, precision=backend_precision, - device=backend_device + device=backend_device, ) - IVPSolver = get_IVP_solver( - library=backend_library + IVPSolver = get_solver_ivp( + library=backend_library, ) # set constants @@ -569,11 +671,11 @@ def __init__(self, # set matrices self.A = backend.zeros( shape=(n_envs, *self.dim_corrs), - dtype='real' + dtype='real', ) self.D = backend.zeros( shape=(n_envs, *self.dim_corrs), - dtype='real' + dtype='real', ) self.is_A_constant = False self.is_D_constant = False @@ -591,9 +693,11 @@ def __init__(self, action_maximums=action_maximums, action_interval=action_interval, data_idxs=data_idxs, - dir_prefix=(dir_prefix if dir_prefix != 'data' else ('data/' + self.name.lower()) + '/env'), - file_prefix='lho_vec_env', - **kwargs + dir_prefix=(dir_prefix \ + if dir_prefix != "data" \ + else ("data/" + self.name.lower()) + "/env"), + file_prefix="lho_vec_env", + **kwargs, ) # initialize solver @@ -606,56 +710,58 @@ def __init__(self, 'atol': kwargs['ode_atol'], 'rtol': kwargs['ode_rtol'], 'is_stiff': False, - 'step_interval': self.action_interval + 'step_interval': self.action_interval, }, func_controls=getattr(self, 'func_controls', None), has_delay=self.has_delay, func_delay=getattr(self, 'func_delay', None), delay_interval=self.action_interval, - backend=self.backend + backend=self.backend, ) # initialize buffers if self.num_modes != 0: self.mode_rates_real = self.backend.zeros( shape=(self.n_envs, 2 * self.num_modes), - dtype='real' + dtype='real', ) if self.num_corrs != 0: self.matmul_0 = self.backend.empty( shape=(self.n_envs, *self.dim_corrs), - dtype='real' + dtype='real', ) self.matmul_1 = self.backend.empty( shape=(self.n_envs, *self.dim_corrs), - dtype='real' + dtype='real', ) self.sum_0 = self.backend.empty( shape=(self.n_envs, *self.dim_corrs), - dtype='real' + dtype='real', ) self.sum_1 = self.backend.empty( shape=(self.n_envs, *self.dim_corrs), - dtype='real' + dtype='real', ) self.y_rates = self.backend.empty( shape=(self.n_envs, 2 * self.num_modes + self.num_corrs), - dtype='real' + dtype='real', ) def _update_states(self): return self.solver.step( T_step=self.T_step, y_0=self.States[-1], - params=self.actions + params=self.actions, ) - def func(self, - t, - y, - args:tuple + def func( + self, + t, + y, + args:tuple, ): - r"""Wrapper function for the rates of change of the real-valued modes and correlations. + r"""Wrapper function for the rates of change + of the real-valued modes and correlations. The variables are cast to real. @@ -664,24 +770,34 @@ def func(self, t: float Time at which the values are calculated. y: Any - Real-valued modes and flattened correlations with shape ``(2 * num_modes + num_corrs, )``. First ``num_modes`` elements contain the real parts of the modes, the next ``num_modes`` elements contain the imaginary parts of the modes, and the last ``num_corrs`` elements contain the correlations. When ``num_modes`` is ``0``, only the correlations are included. When ``num_corrs`` is ``0``, only the modes are included. + Real-valued modes and flattened correlations + with shape ``(2 * num_modes + num_corrs, )``. + First ``num_modes`` elements contain the real parts of the modes, + the next ``num_modes`` elements contain the imaginary parts + of the modes, and the last ``num_corrs`` elements + contain the correlations. + When ``num_modes`` is ``0``, only the correlations are included. + When ``num_corrs`` is ``0``, only the modes are included. args: tuple Actions, control function and delay function. Returns ------- rates: Any - Rates of change of the real-valued modes and flattened correlations with shape ``(2 * num_modes + num_corrs, )``. + Rates of change of the real-valued modes + and flattened correlations with shape + ``(2 * num_modes + num_corrs, )``. """ # extract frequently used variables if self.num_modes != 0: - modes = y[:, :self.num_modes] + 1.0j * y[:, self.num_modes:2 * self.num_modes] + modes = y[:, :self.num_modes] \ + + 1.0j * y[:, self.num_modes:2 * self.num_modes] # get real-valued mode rates _mode_rates_real = self.get_mode_rates_real( t=t, modes_real=y[:, :2 * self.num_modes], - args=args + args=args, ) if self.num_corrs == 0: return _mode_rates_real @@ -691,21 +807,21 @@ def func(self, if self.num_corrs != 0: corrs = self.backend.reshape( tensor=y[:, 2 * self.num_modes:], - shape=(self.n_envs, *self.dim_corrs) + shape=(self.n_envs, *self.dim_corrs), ) # get drift matrix A = self.A if self.is_A_constant else self.get_A( t=t, modes=modes, - args=args + args=args, ) # get noise matrix D = self.D if self.is_D_constant else self.get_D( t=t, modes=modes, - args=args + args=args, ) # get flattened correlation rates @@ -715,7 +831,7 @@ def func(self, tensor_0=self.backend.matmul( tensor_0=A, tensor_1=corrs, - out=self.matmul_0 + out=self.matmul_0, ), tensor_1=self.backend.matmul( tensor_0=corrs, @@ -724,13 +840,13 @@ def func(self, axis_0=1, axis_1=2 ), - out=self.matmul_1 + out=self.matmul_1, ), out=self.sum_0), tensor_1=D, - out=self.sum_1 + out=self.sum_1, ), - shape=(self.n_envs, self.num_corrs) + shape=(self.n_envs, self.num_corrs), ) if self.num_modes == 0: @@ -739,16 +855,17 @@ def func(self, return self.backend.concatenate( tensors=( _mode_rates_real, - _corr_rates_flat + _corr_rates_flat, ), axis=1, - out=self.y_rates + out=self.y_rates, ) - def get_A(self, - t, - modes, - args:tuple + def get_A( + self, + t, + modes, + args:tuple, ): """Method to obtain the Jacobian of quantum fluctuation quadratures. @@ -757,22 +874,25 @@ def get_A(self, t: float Time at which the values are calculated. modes: Any - Classical mode amplitudes with shape ``(n_envs, num_modes)``. + Classical mode amplitudes + with shape ``(n_envs, num_modes)``. args: tuple Actions, control function and delay function. Returns ------- A: Any - Jacobian of quantum fluctuation quadratures with shape ``(n_envs, num_quads, num_quads)``. + Jacobian of quantum fluctuation quadratures + with shape ``(n_envs, num_quads, num_quads)``. """ raise NotImplementedError - def get_D(self, - t, - modes, - args:tuple + def get_D( + self, + t, + modes, + args:tuple, ): """Method to obtain the quantum noise correaltions. @@ -781,64 +901,75 @@ def get_D(self, t: float Time at which the values are calculated. modes: Any - Classical mode amplitudes with shape ``(n_envs, num_modes)``. + Classical mode amplitudes + with shape ``(n_envs, num_modes)``. args: tuple Actions, control function and delay function. Returns ------- D: Any - Quantum noise correlations with shape ``(n_envs, num_quads, num_quads)``. + Quantum noise correlations with shape + ``(n_envs, num_quads, num_quads)``. """ raise NotImplementedError - def get_mode_rates(self, - t, - modes, - args:tuple + def get_mode_rates( + self, + t, + modes, + args:tuple, ): - """Method to obtain the rates of change of the classical mode amplitudes. + """Method to obtain the rates of change + of the classical mode amplitudes. Parameters ---------- t: float Time at which the values are calculated. modes: Any - Classical mode amplitudes with shape ``(n_envs, num_modes)``. + Classical mode amplitudes + with shape ``(n_envs, num_modes)``. args: tuple Actions, control function and delay function. Returns ------- D: Any - Rates of change of the classical mode amplitudes with shape ``(n_envs, num_modes)``. + Rates of change of the classical mode amplitudes + with shape ``(n_envs, num_modes)``. """ raise NotImplementedError - def get_mode_rates_real(self, - t, - modes_real, - args:tuple + def get_mode_rates_real( + self, + t, + modes_real, + args:tuple, ): - """Method to obtain the real-valued mode rates from real-valued modes. + """Method to obtain the real-valued mode rates + from real-valued modes. - The interfaced environment needs to implement the ``get_mode_rates`` method. + The interfaced environment needs + to implement the ``get_mode_rates`` method. Parameters ---------- t: float Time at which the values are calculated. modes_real: Any - Real-valued classical mode amplitudes with shape ``(n_envs, 2 * num_modes)``. + Real-valued classical mode amplitudes + with shape ``(n_envs, 2 * num_modes)``. args: tuple Actions, control function and delay function. Returns ------- mode_rates_real: Any - Real-valued rates of change of the classcial mode amplitudes with shape ``(n_envs, 2 * num_modes)``. + Real-valued rates of change of the classcial mode + amplitudes with shape ``(n_envs, 2 * num_modes)``. """ # handle null @@ -848,20 +979,21 @@ def get_mode_rates_real(self, # get complex-valued mode rates mode_rates = self.get_mode_rates( t=t, - modes=modes_real[:, :self.num_modes] + 1.0j * modes_real[:, self.num_modes:], - args=args + modes=modes_real[:, :self.num_modes] \ + + 1.0j * modes_real[:, self.num_modes:], + args=args, ) # return real-valued mode rates return self.backend.concatenate( tensors=( self.backend.real( - tensor=mode_rates + tensor=mode_rates, ), self.backend.imag( - tensor=mode_rates + tensor=mode_rates, ) ), axis=1, - out=self.mode_rates_real + out=self.mode_rates_real, ) diff --git a/quantrl/envs/stochastic.py b/quantrl/envs/stochastic.py index 42c6c0f..b31c312 100644 --- a/quantrl/envs/stochastic.py +++ b/quantrl/envs/stochastic.py @@ -6,29 +6,37 @@ __name__ = 'quantrl.envs.stochastic' __authors__ = ["Sampreet Kalita"] __created__ = "2023-04-25" -__updated__ = "2025-05-11" +__updated__ = "2025-08-20" # dependencies import numpy as np # quantrl modules -from ..backends.context_manager import get_backend_instance +from ..backends.context_manager import get_instance_backend from .base import BaseGymEnv # TODO: Add delay feature # TODO: Release memory class LinearEnv(BaseGymEnv): - """Class to interface stochastic linear environments using Wiener increments. + """Class to interface stochastic linear + environments using Wiener increments. Initializes ``A`` and ``is_A_constant``. - The interfaced environment requires ``default_params`` dictionary defined before initializing the parent class. + The interfaced environment requires + ``default_params`` dictionary defined + before initializing the parent class. - The interfaced environment needs to implement ``reset_states`` and ``get_reward`` methods. - Additionally, the ``get_properties`` method should be overridden if ``n_properties`` is non-zero. - Refer to **Notes** of :class:`quantrl.envs.base.BaseEnv` for their implementations. + The interfaced environment needs to implement + ``reset_states`` and ``get_reward`` methods. + Additionally, the ``get_properties`` method + should be overridden if ``n_properties`` is non-zero. + Refer to **Notes** of :class:`quantrl.envs.base.BaseEnv` + for their implementations. - The ``func`` method requires ``get_A`` for the Jacobian of the states and the ``get_noise_prefixes`` for the noise values. + The ``func`` method requires ``get_A`` + for the Jacobian of the states and the + ``get_noise_prefixes`` for the noise values. Parameters ---------- @@ -53,19 +61,30 @@ class LinearEnv(BaseGymEnv): action_maximums: list Maximum values of each action. action_interval: int - Interval at which the actions are updated. Must be positive. + Interval at which the actions are updated. + Must be positive. data_idxs: list - Indices of the data to store into the ``data`` attribute. The indices can be selected from the complete set of values at each point of time (total ``1 + n_actions + n_observations + n_properties + 1`` elements in the same order, where the first element is the time and the last element is the reward). + Indices of the data to store into the ``data`` attribute. + The indices can be selected from the + complete set of values at each point of time + (total ``1 + n_actions + n_observations + n_properties + 1`` + elements in the same order, where the first element + is the time and the last element is the reward). backend_library: str, default='numpy' - Solver to use for each step. Options are ``'jax'`` for JAX-based solvers, ``'torch'`` for PyTorch-based solvers and ``'numpy'`` for NumPy/SciPy-based solvers. + Solver to use for each step. + Options are ``'jax'`` for JAX-based solvers, + ``'torch'`` for PyTorch-based solvers and + ``'numpy'`` for NumPy/SciPy-based solvers. backend_precision: str, default='double' - Precision of the numerical values in the backend. Options are ``'single'`` and ``'double'``. + Precision of the numerical values in the backend. + Options are ``'single'`` and ``'double'``. backend_device: str, default='cuda' Device to run the solver. Options are ``'cpu'`` and ``'cuda'``. - dir_prefix: str, default='data' + dir_prefix: str, default="data" Prefix of the directory where the data will be stored. kwargs: dict, optional - Keyword arguments. Refer to the ``kwargs`` parameter of :class:`quantrl.envs.base.BaseEnv` for available options. + Keyword arguments. Refer to the ``kwargs`` parameter of + :class:`quantrl.envs.base.BaseEnv` for available options. """ default_params = {} @@ -74,35 +93,37 @@ class LinearEnv(BaseGymEnv): backend_libraries = ['jax', 'torch', 'numpy'] """list: Available backend libraries.""" - def __init__(self, - name:str, - desc:str, - params:dict, - t_norm_max:float, - t_norm_ssz:float, - t_norm_mul:float, - n_observations:int, - n_properties:int, - n_actions:int, - action_maximums:list, - action_interval:int, - data_idxs:list, - backend_library:str='numpy', - backend_precision:str='double', - backend_device:str='cuda', - dir_prefix:str='data', - **kwargs + def __init__( + self, + name:str, + desc:str, + params:dict, + t_norm_max:float, + t_norm_ssz:float, + t_norm_mul:float, + n_observations:int, + n_properties:int, + n_actions:int, + action_maximums:list, + action_interval:int, + data_idxs:list, + backend_library:str='numpy', + backend_precision:str='double', + backend_device:str='cuda', + dir_prefix:str="data", + **kwargs, ): """Class constructor for LinearEnv.""" # validate arguments - assert backend_library in self.backend_libraries, f"parameter ``backend_library`` should be one of ``{self.backend_libraries}``" + assert backend_library in self.backend_libraries, \ + f"parameter ``backend_library`` should be one of ``{self.backend_libraries}``" # select backend - backend = get_backend_instance( + backend = get_instance_backend( library=backend_library, precision=backend_precision, - device=backend_device + device=backend_device, ) # set constants @@ -117,11 +138,11 @@ def __init__(self, self.Ws = None self.I = backend.eye( N=n_observations, - dtype='real' + dtype='real', ) self.A = backend.zeros( shape=(n_observations, n_observations), - dtype='real' + dtype='real', ) self.is_A_constant = False @@ -137,42 +158,45 @@ def __init__(self, action_maximums=action_maximums, action_interval=action_interval, data_idxs=data_idxs, - dir_prefix=(dir_prefix if dir_prefix != 'data' else ('data/' + self.name.lower()) + '/env') + '_' + '_'.join([ - str(val) for _, val in self.params.items() - ]), - file_prefix='lin_env', - **kwargs + dir_prefix=(dir_prefix \ + if dir_prefix != "data" \ + else ("data/" + self.name.lower()) + "/env") + "_" + "_".join( + [str(val) for _, val in self.params.items()] + ), + file_prefix="lin_env", + **kwargs, ) # initialize buffers self.add_0 = self.backend.empty( shape=(self.n_observations, ), - dtype='real' + dtype='real', ) self.matmul_0 = self.backend.empty( shape=(self.n_observations, ), - dtype='real' + dtype='real', ) - def reset(self, - seed:float=None, - options:dict=None + def reset( + self, + seed:float=None, + options:dict=None, ): # update Wiener noises self.Ws = np.sqrt(self.t_ssz) * self.backend.normal( generator=self.backend.generator( - seed=self.seed + seed=self.seed, ), shape=(self.shape_T[0] - 1, self.n_observations), mean=0.0, std=1.0, - dtype='real' + dtype='real', ) # return observations return super().reset( seed=seed, - options=options + options=options, ) def _update_states(self): @@ -180,25 +204,27 @@ def _update_states(self): _States = self.backend.jit_update( tensor=self.States, indices=0, - values=self.States[-1] + values=self.States[-1], ) # iterate and return return self.backend.iterate_i( func=self.func, iterations_i=self.backend.shape( - tensor=self.T_step + tensor=self.T_step, )[0] - 1, Y=_States, - args=(self.actions, None, None) + args=(self.actions, None, None), ) - def func(self, - i, - Y, - args:tuple + def func( + self, + i, + Y, + args:tuple, ): - """Method to obtain the rates of change of the real-valued variables. + """Method to obtain the rates of change + of the real-valued variables. Parameters ---------- @@ -212,38 +238,41 @@ def func(self, Returns ------- rates: Any - Rates of change of the real-valued modes and flattened correlations with shape ``(2 * num_modes + num_corrs, )``. + Rates of change of the real-valued modes + and flattened correlations with shape + ``(2 * num_modes + num_corrs, )``. """ # get drift matrix M_i = self.I + self.get_A( t_idx=self.t_idx + i, - args=args + args=args, ) * self.t_ssz # get noise prefixes n_i = self.get_noise_prefixes( t_idx=self.t_idx + i, - args=args + args=args, ) # get updated states values = self.backend.jit_add( tensor_0=self.backend.jit_matmul( tensor_0=M_i, tensor_1=Y[i], - out=self.matmul_0 + out=self.matmul_0, ), tensor_1=n_i * self.Ws[self.t_idx + i], - out=self.add_0 + out=self.add_0, ) return self.backend.jit_update( tensor=Y, indices=i + 1, - values=values + values=values, ) - def get_A(self, - t_idx:int, - args:tuple + def get_A( + self, + t_idx:int, + args:tuple, ): """Method to obtain the Jacobian of the states. @@ -257,14 +286,16 @@ def get_A(self, Returns ------- A: Any - Jacobian of the states with shape ``(n_observations, n_observations)``. + Jacobian of the states with shape + ``(n_observations, n_observations)``. """ raise NotImplementedError - def get_noise_prefixes(self, - t_idx:int, - args:tuple + def get_noise_prefixes( + self, + t_idx:int, + args:tuple, ): """Method to obtain the noise prefixes for each state. @@ -278,7 +309,8 @@ def get_noise_prefixes(self, Returns ------- noise_prefixes: Any - Noise prefixes for each observation with shape ``(n_observations, )``. + Noise prefixes for each observation with shape + ``(n_observations, )``. """ raise NotImplementedError diff --git a/quantrl/io.py b/quantrl/io.py index db9e94a..aae85a5 100644 --- a/quantrl/io.py +++ b/quantrl/io.py @@ -6,7 +6,7 @@ __name__ = 'quantrl.io' __authors__ = ["Sampreet Kalita"] __created__ = "2023-12-07" -__updated__ = "2025-04-21" +__updated__ = "2025-08-20" # dependencies import gc @@ -23,25 +23,34 @@ class FileIO(): """Handler for file input-output. Initializes ``cache`` to ``None`` and ``index`` to ``-1``. - Subsequent calls to ``update_cache`` allocates ``cache`` and updates ``index``. - The parent needs to implement the ``close`` method to cache the final file. + Subsequent calls to ``update_cache`` allocates + ``cache`` and updates ``index``. + The parent needs to implement the ``close`` method + to cache the final file. Parameters ---------- disk_cache_dir: str - Directory path for the disk cache. If the value of ``disk_cache_size`` is ``0``, the then this parameter serves as the file path for a single disk cache, else the cache is dumped in parts. + Directory path for the disk cache. + If the value of ``disk_cache_size`` is ``0``, + then this parameter serves as the file path for a single disk cache, + else the cache is dumped in parts. cache_dump_interval: int, default=100 - Number of steps to update the cache before dumping it to disk. Should be a positive integer. + Number of steps to update the cache before dumping it to disk. + Should be a positive integer. """ - def __init__(self, - disk_cache_dir:str, - cache_dump_interval:int=100 + def __init__( + self, + disk_cache_dir:str, + cache_dump_interval:int=100, ): """Class constructor for FileIO.""" # set attributes - assert isinstance(cache_dump_interval, int) and cache_dump_interval > 0, "parameter ``disk_cache_size`` should be a positive integer" + assert isinstance(cache_dump_interval, int) \ + and cache_dump_interval > 0, \ + "parameter ``disk_cache_size`` should be a positive integer" self.disk_cache_dir = disk_cache_dir self.cache_dump_interval = cache_dump_interval try: @@ -53,10 +62,11 @@ def __init__(self, self.cache = None self.index = -1 - def dump_part_async(self, - data:np.ndarray, - batch_idx:int, - part_idx:int + def dump_part_async( + self, + data:np.ndarray, + batch_idx:int, + part_idx:int, ): """Method to dump a batch of data to disk asynchronously. @@ -71,11 +81,14 @@ def dump_part_async(self, """ # save as compressed NumPy data from another thread - thread = Thread(target=np.savez_compressed, args=(self.disk_cache_dir + '/' + '_'.join([ - str(batch_idx * self.cache_dump_interval), - str((batch_idx + 1) * self.cache_dump_interval - 1), - str(part_idx) - ]) + '.npz', data)) + thread = Thread( + target=np.savez_compressed, + args=(self.disk_cache_dir + "/" + "_".join([ + str(batch_idx * self.cache_dump_interval), + str((batch_idx + 1) * self.cache_dump_interval - 1), + str(part_idx), + ]) + ".npz", data), + ) thread.start() thread.join() @@ -83,8 +96,9 @@ def dump_part_async(self, del data gc.collect() - def update_cache(self, - data:np.ndarray + def update_cache( + self, + data:np.ndarray, ): """Method to update the cache with data. @@ -96,21 +110,26 @@ def update_cache(self, # update list if self.cache is None: - self.cache = np.zeros((self.cache_dump_interval, *data.shape), dtype=data.dtype) + self.cache = np.zeros(( + self.cache_dump_interval, + *data.shape, + ), dtype=data.dtype) self.index += 1 self.cache[self.index % self.cache_dump_interval] = data # dump cache - if self.index != 0 and (self.index + 1) % self.cache_dump_interval == 0: + if self.index != 0 \ + and (self.index + 1) % self.cache_dump_interval == 0: self._dump_cache_async( idx_start=self.index - self.cache_dump_interval + 1, - idx_end=self.index + idx_end=self.index, ) - def _dump_cache_async(self, - idx_start:int, - idx_end:int + def _dump_cache_async( + self, + idx_start:int, + idx_end:int, ): """Method to dump cache to disk asynchronously. @@ -123,7 +142,11 @@ def _dump_cache_async(self, """ # save as compressed NumPy data from another thread - thread = Thread(target=np.savez_compressed, args=(self.disk_cache_dir + '/' + str(idx_start) + '_' + str(idx_end) + '.npz', self.cache)) + thread = Thread( + target=np.savez_compressed, + args=(self.disk_cache_dir + "/" + str(idx_start) \ + + "_" + str(idx_end) + ".npz", self.cache), + ) thread.start() # clear cache @@ -131,12 +154,14 @@ def _dump_cache_async(self, self.cache = None gc.collect() - def get_disk_cache(self, - idx_start:int=0, - idx_end:int=-1, - idxs:list=None + def get_disk_cache( + self, + idx_start:int=0, + idx_end:int=-1, + idxs:list=None, ): - """Method to return select disk-cached data between a given set of indices. + """Method to return select disk-cached data + between a given set of indices. Parameters ---------- @@ -145,35 +170,44 @@ def get_disk_cache(self, idx_end: int, default=-1 Ending index for the part file. idxs: list or slice, default=None - Indices of the data values required. If ``None``, all data is returned. + Indices of the data values required. + If ``None``, all data is returned. """ # iterate over parts cache_list = [] for i in tqdm( - range(int(idx_start / self.cache_dump_interval) * self.cache_dump_interval, idx_end + 1, self.cache_dump_interval), + range( + int(idx_start / self.cache_dump_interval) \ + * self.cache_dump_interval, idx_end + 1, + self.cache_dump_interval, + ), desc="Loading", leave=False, mininterval=0.5, - disable=False + disable=False, ): # update end index _idx_e = i + self.cache_dump_interval - 1 # update cache list _cache = self._load_cache( idx_start=i, - idx_end=_idx_e + idx_end=_idx_e, ) - cache_list += [_cache[:, :, idxs].copy() if idxs is not None else _cache.copy()] + cache_list += [_cache[:, :, idxs].copy() \ + if idxs is not None \ + else _cache.copy()] # clear loaded cache del _cache gc.collect() - return np.concatenate(cache_list)[idx_start % self.cache_dump_interval:] + return np.concatenate(cache_list)[idx_start \ + % self.cache_dump_interval:] - def _load_cache(self, - idx_start:int, - idx_end:int + def _load_cache( + self, + idx_start:int, + idx_end:int, ): """Method to load cache from disk. @@ -186,11 +220,18 @@ def _load_cache(self, """ # load part or single cache file - return np.load(self.disk_cache_dir + '/' + str(idx_start) + '_' + (str(idx_end) if idx_end != -1 else '*') + '.npz')['arr_0'] - - def save_data(self, - data:np.ndarray, - file_name:str + return np.load( + self.disk_cache_dir \ + + "/" + str(idx_start) \ + + "_" + (str(idx_end) \ + if idx_end != -1 \ + else "*") + ".npz", + )['arr_0'] + + def save_data( + self, + data:np.ndarray, + file_name:str, ): """Method to save data to a file. @@ -202,10 +243,11 @@ def save_data(self, Name of the file. """ - np.savez_compressed(file_name + '.npz', data) + np.savez_compressed(file_name + ".npz", data) - def load_data(self, - file_name:str + def load_data( + self, + file_name:str, ): """Method to load data from a file. @@ -217,15 +259,17 @@ def load_data(self, Returns ------- data: :class:`numpy.ndarray` - Data loaded from the file. Returns `None` if the file does not exist. + Data loaded from the file. + Returns `None` if the file does not exist. """ - if os.path.isfile(file_name + '.npz'): - return np.load(file_name + '.npz')['arr_0'] + if os.path.isfile(file_name + ".npz"): + return np.load(file_name + ".npz")['arr_0'] return None - def close(self, - dump_cache=True + def close( + self, + dump_cache=True, ): """Method to close FileIO. @@ -236,10 +280,11 @@ def close(self, """ if dump_cache and self.cache is not None: - _idx_s = self.index - (self.index + 1) % self.cache_dump_interval + 1 + _idx_s = self.index - (self.index + 1) \ + % self.cache_dump_interval + 1 self._dump_cache_async( idx_start=_idx_s, - idx_end=_idx_s + self.cache_dump_interval - 1 + idx_end=_idx_s + self.cache_dump_interval - 1, ) # clean diff --git a/quantrl/plotters.py b/quantrl/plotters.py index e78e0d5..d49495a 100644 --- a/quantrl/plotters.py +++ b/quantrl/plotters.py @@ -6,7 +6,7 @@ __name__ = 'quantrl.plotters' __authors__ = ["Sampreet Kalita"] __created__ = "2023-12-08" -__updated__ = "2024-10-14" +__updated__ = "2025-08-20" # dependencies from io import BytesIO @@ -44,29 +44,39 @@ class TrajectoryPlotter(): Parameters ---------- axes_args: list - Lists of axis properties. The first element of each entry is the ``x_label``, the second is ``y_label``, the third is ``[y_limit_min, y_limit_max]`` and the fourth is ``y_scale``. + Lists of axis properties. + The first element of each entry is the ``x_label``, + the second is ``y_label``, + the third is ``[y_limit_min, y_limit_max]`` and + the fourth is ``y_scale``. axes_lines_max: int, default=10 - Maximum number of lines to display in each plot. Higher numbers slow down the run. Default is ``10``. + Maximum number of lines to display in each plot. + Higher numbers slow down the run. Default is ``10``. axes_cols: int, default=2 - Number of columns in the figure. Default is ``2``. + Number of columns in the figure. + Default is ``2``. show_title: bool, default=True Option to display the trajectory index as title. save_dir: str, default=None - Directory to save the plots on each update. If ``None``, the plots are not saved. + Directory to save the plots on each update. + If ``None``, the plots are not saved. """ - def __init__(self, - axes_args:list, - axes_lines_max:int=10, - axes_cols:int=2, - show_title:bool=True, - save_dir:str=None + def __init__( + self, + axes_args:list, + axes_lines_max:int=10, + axes_cols:int=2, + show_title:bool=True, + save_dir:str=None, ): """Class constructor for TrajectoryPlotter.""" # validate - assert axes_lines_max >= 0, "parameter ``axes_lines_max`` should be a non-negative integer" - assert axes_cols > 0, "parameter ``axes_cols`` should be a positive integer" + assert axes_lines_max >= 0, \ + "parameter ``axes_lines_max`` should be a non-negative integer" + assert axes_cols > 0, \ + "parameter ``axes_cols`` should be a positive integer" # set attributes self.axes_args = axes_args @@ -87,8 +97,16 @@ def __init__(self, # initialize variables self.axes_rows = int(np.ceil(len(self.axes_args) / self.axes_cols)) - self.fig = plt.figure(figsize=(6.0 * self.axes_cols, 3.0 * self.axes_rows)) - self.gspec = GridSpec(self.axes_rows, self.axes_cols, figure=self.fig, width_ratios=[0.2] * self.axes_cols) + self.fig = plt.figure(figsize=( + 6.0 * self.axes_cols, + 3.0 * self.axes_rows, + )) + self.gspec = GridSpec( + self.axes_rows, + self.axes_cols, + figure=self.fig, + width_ratios=[0.2] * self.axes_cols, + ) self.axes = [] self.lines = None @@ -109,17 +127,18 @@ def __init__(self, ax.set_yscale(ax_args[3]) self.axes.append(ax) if self.show_title: - self.fig.suptitle('#0') + self.fig.suptitle("#0") self.fig.tight_layout() # initialize buffers self.frames = [] - def plot_lines(self, - xs, - Y, - traj_idx=0, - update_buffer=False + def plot_lines( + self, + xs, + Y, + traj_idx=0, + update_buffer=False, ): """Method to plot new lines. @@ -143,12 +162,13 @@ def plot_lines(self, # add new lines self.lines = [] for i, ax in enumerate(self.axes): - if self.axes_lines_max and len(ax.get_lines()) >= self.axes_lines_max: + if self.axes_lines_max \ + and len(ax.get_lines()) >= self.axes_lines_max: line = ax.get_lines()[0] line.remove() self.lines.append(ax.plot(xs, Y[:, i])[0]) if self.show_title: - self.fig.suptitle('#' + str(traj_idx)) + self.fig.suptitle("#" + str(traj_idx)) self.fig.canvas.draw() self.fig.canvas.flush_events() @@ -161,14 +181,15 @@ def plot_lines(self, # save plot if self.save_dir is not None: self.save_plot( - file_name=self.save_dir + '/traj_' + str(traj_idx) + file_name=self.save_dir + "/traj_" + str(traj_idx), ) self.save_plot( - file_name=self.save_dir + '_latest' + file_name=self.save_dir + "_latest", ) - def make_gif(self, - file_name:str + def make_gif( + self, + file_name:str, ): """Method to create a gif file from the frame buffer. @@ -184,14 +205,22 @@ def make_gif(self, # dump buffer frame = self.frames[0] - frame.save(file_name + '.gif', format='GIF', append_images=self.frames[1:], save_all=True, duration=50, loop=0) + frame.save( + file_name + ".gif", + format='GIF', + append_images=self.frames[1:], + save_all=True, + duration=50, + loop=0, + ) # reset buffer del self.frames self.frames = [] - def save_plot(self, - file_name:str + def save_plot( + self, + file_name:str, ): """Method to save the plot. @@ -232,24 +261,32 @@ class LearningCurvePlotter(): Parameters ---------- axes_args: list - Lists of axis properties. The first element of each entry is the ``x_label``, the second is ``y_label``, the third is ``[y_limit_min, y_limit_max]`` and the fourth is ``y_scale``. + Lists of axis properties. + The first element of each entry is the ``x_label``, + the second is ``y_label``, + the third is ``[y_limit_min, y_limit_max]`` and + the fourth is ``y_scale``. average_over: int, default=100 Number of points to average over. percentiles: list, default=None - Percentile values for intraquartile ranges. If ``None``, the percentiles are set to ``[25, 50, 75]``. + Percentile values for intraquartile ranges. + If ``None``, the percentiles are set to ``[25, 50, 75]``. """ - def __init__(self, - axis_args:list, - average_over:int=100, - percentiles:list=None + def __init__( + self, + axis_args:list, + average_over:int=100, + percentiles:list=None, ): """Class constructor for LearningCurvePlotter.""" # set attributes self.axis_args = axis_args self.average_over = average_over - self.percentiles = percentiles if percentiles is not None else [25, 50, 75] + self.percentiles = percentiles \ + if percentiles is not None \ + else [25, 50, 75] # turn on interactive mode plt.ion() @@ -268,12 +305,13 @@ def __init__(self, self.line = None self.line_faint = None - def add_data(self, - data_rewards:np.ndarray, - renew:bool=False, - color:str='k', - style:str='-', - width:float=1.5 + def add_data( + self, + data_rewards:np.ndarray, + renew:bool=False, + color:str='k', + style:str='-', + width:float=1.5, ): """Method to add reward data. @@ -294,7 +332,11 @@ def add_data(self, # if averaging opted data_rewards_smooth = data_rewards if self.average_over is not None: - data_rewards_smooth = np.convolve(data_rewards, np.ones((self.average_over, )) / float(self.average_over), mode='valid') + data_rewards_smooth = np.convolve( + data_rewards, + np.ones((self.average_over, )) / float(self.average_over), + mode='valid', + ) # update data if renew: @@ -306,27 +348,59 @@ def add_data(self, self.line_faint.remove() if self.line is not None: self.line.remove() - xs = list(range(self.average_over, len(self.data[0]) + self.average_over)) + xs = list(range( + self.average_over, + len(self.data[0]) + self.average_over, + )) # if single entry if len(self.data) == 1: q_mean = data_rewards_smooth - self.line_faint = self.ax.plot(xs, data_rewards[self.average_over - 1:], c=color, alpha=0.1, linewidth=0.5)[0] + self.line_faint = self.ax.plot( + xs, + data_rewards[self.average_over - 1:], + c=color, + alpha=0.1, + linewidth=0.5, + )[0] else: # if interquartile ranges given if self.percentiles is not None: - q_min, q_mean, q_max = np.percentile(self.data, self.percentiles, 0) - self.line_faint = self.ax.fill_between(xs, q_min, q_max, facecolor=color, alpha=0.1) + q_min, q_mean, q_max = np.percentile( + self.data, + self.percentiles, + 0, + ) + self.line_faint = self.ax.fill_between( + xs, + q_min, + q_max, + facecolor=color, + alpha=0.1, + ) # calculate mean else: q_mean = np.mean(self.data, 0) - self.line_faint = self.ax.plot(xs, q_mean, c=color, alpha=0.1, linewidth=0.5)[0] + self.line_faint = self.ax.plot( + xs, + q_mean, + c=color, + alpha=0.1, + linewidth=0.5, + )[0] # add new lines - self.line = self.ax.plot(xs, q_mean, c=color, linestyle=style, linewidth=width)[0] - - def save_plot(self, - file_name:str + self.line = self.ax.plot( + xs, + q_mean, + c=color, + linestyle=style, + linewidth=width, + )[0] + + def save_plot( + self, + file_name:str, ): """Method to save the plot. diff --git a/quantrl/solvers/base.py b/quantrl/solvers/base.py index 464d5a5..f8007be 100644 --- a/quantrl/solvers/base.py +++ b/quantrl/solvers/base.py @@ -6,7 +6,7 @@ __name__ = 'quantrl.solvers.base' __authors__ = ["Sampreet Kalita"] __created__ = "2024-03-10" -__updated__ = "2025-05-12" +__updated__ = "2025-08-20" # dependencies from abc import ABC, abstractmethod @@ -21,7 +21,8 @@ class BaseIVPSolver(ABC): """ODE and DDE solver backend for initial-value problems. - The inherited classes should contain the ``solver_methods`` attribute listing all the available methods of the corresponding solver. + The inherited classes should contain the ``solver_methods`` attribute + listing all the available methods of the corresponding solver. Currently, the module only supports a single delay interval. @@ -29,31 +30,45 @@ class BaseIVPSolver(ABC): ---------- func: callable ODE/DDE function in the format ``func(t, y, args)``. - The first element of ``args`` contains the constant parameters, the second element contains the function for the controls and the third contains the delay function. + The first element of ``args`` contains the constant parameters, + the second element contains the function for the controls + and the third contains the delay function. solver_params: dict Parameters of the solver. Currently supported options are: - ================ ==================================================== + ================ ============================================ key value - ================ ==================================================== - method (*str*) method used to solve the ODEs. Refer to the documentation of the inherited solvers. Default is ``'vode'`` from :class:`quantrl.solvers.numpy.SciPyIVPSolver`. - atol (*float*) absolute tolerance of the integrator. Default is ``1e-12``. - rtol (*float*) relative tolerance of the integrator. Default is ``1e-9``. - is_stiff (*bool*) option to select whether the integration is a stiff problem or a non-stiff one. Default is ``False``. - step_interval (*bool*) number of steps to jump during integration. Higher values give faster results. Default is ``10``. - ================ ==================================================== + ================ ============================================ + method (*str*) method used to solve the ODEs. + Refer to the documentation of the inherited + solvers. Default is ``'vode'`` from + :class:`quantrl.solvers.numpy.SciPyIVPSolver`. + atol (*float*) absolute tolerance of the integrator. + Default is ``1e-12``. + rtol (*float*) relative tolerance of the integrator. + Default is ``1e-9``. + is_stiff (*bool*) option to select whether the + integration is a stiff problem or + a non-stiff one. Default is ``False``. + step_interval (*bool*) number of steps to jump during + integration. Higher values give faster results. + Default is ``10``. + ================ ============================================ func_controls: callable Function for the controls in the format ``func_controls(t)``. has_delay: bool Option to solve DDEs. func_delay: callable History function for first delay step in the format ``func_delay(t)``. - This function is then internally replaced by the interpolated function for the subsequent steps. + This function is then internally replaced + by the interpolated function for the subsequent steps. delay_interval: int Interval of the delay. - ..note: In the presence of delay, the parameter ``'step_interval'`` is overriden by the delay interval. + ..note: In the presence of delay, + the parameter ``'step_interval'`` + is overriden by the delay interval. """ # attributes @@ -63,22 +78,23 @@ class BaseIVPSolver(ABC): 'rtol': 1e-9, 'is_stiff': False, 'step_interval': 10, - 'complex': False + 'complex': False, } """dict: Default parameters of the solver.""" solver_methods = [] """list: Methods used by the solver.""" - def __init__(self, - func, - y_0, - T, - solver_params:dict, - func_controls, - has_delay:bool, - func_delay, - delay_interval:int, - backend:BaseBackend + def __init__( + self, + func, + y_0, + T, + solver_params:dict, + func_controls, + has_delay:bool, + func_delay, + delay_interval:int, + backend:BaseBackend, ): """Class constructor for BaseIVPSolver.""" @@ -93,36 +109,47 @@ def __init__(self, self.backend = backend # validate attributes - assert self.func_delay is not None if self.has_delay else True, "delay function cannot be ``None`` if parameter ``has_delay`` is ``True``" + assert self.func_delay is not None \ + if self.has_delay \ + else True, \ + "delay function cannot be ``None`` if parameter ``has_delay`` is ``True``" # frequently used variables self.shape_y = self.backend.shape( - tensor=self.y_0 + tensor=self.y_0, ) self.shape_T = self.backend.shape( - tensor=self.T + tensor=self.T, ) # set params self.solver_params = {} for key, _ in self.default_solver_params.items(): - self.solver_params[key] = solver_params.get(key, self.default_solver_params[key]) + self.solver_params[key] = solver_params.get( + key, + self.default_solver_params[key] + ) # override step dimension with delay interval if DDE if self.has_delay and self.delay_interval != 0: self.solver_params['step_interval'] = self.delay_interval # validate params - assert self.solver_params['method'] in self.solver_methods, f"parameter ``method`` should be one of ``{self.solver_methods}``" - assert isinstance(self.solver_params['step_interval'], int) and self.solver_params['step_interval'] < self.shape_T[0], "parameter ``step_interval`` should be an integer with a value less than the total number of steps" + assert self.solver_params['method'] in self.solver_methods, \ + f"parameter ``method`` should be one of ``{self.solver_methods}``" + assert isinstance(self.solver_params['step_interval'], int) \ + and self.solver_params['step_interval'] < self.shape_T[0], \ + "parameter ``step_interval`` should be an integer " \ + + "with a value less than the total number of steps" # step constants self.step_interval = self.solver_params['step_interval'] @abstractmethod - def integrate(self, - T_step, - y_0, - params + def integrate( + self, + T_step, + y_0, + params, ): """Method to take one integration step. @@ -144,9 +171,10 @@ def integrate(self, raise NotImplementedError @abstractmethod - def interpolate(self, - T_step, - Y + def interpolate( + self, + T_step, + Y, ): """Method to take one interpolation step. @@ -165,10 +193,11 @@ def interpolate(self, raise NotImplementedError - def step(self, - T_step, - y_0, - params + def step( + self, + T_step, + y_0, + params, ): """Method to take one integration and interpolation step. @@ -191,22 +220,23 @@ def step(self, _Y = self.integrate( T_step=T_step, y_0=y_0, - params=params + params=params, ) # update delay function if self.has_delay: self.func_delay = self.interpolate( T_step=T_step, - Y=_Y + Y=_Y, ) return _Y - def solve_ivp(self, - y_0, - params, - show_progress + def solve_ivp( + self, + y_0, + params, + show_progress, ): """Module to solve the IVP. @@ -229,13 +259,13 @@ def solve_ivp(self, Y = self.backend.empty( shape=( *self.backend.shape( - tensor=self.T + tensor=self.T, ), *self.backend.shape( - tensor=y_0 + tensor=y_0, ) ), - dtype='real' + dtype='real', ) Y[0] = y_0 @@ -245,12 +275,12 @@ def solve_ivp(self, desc="Solving", leave=False, mininterval=0.5, - disable=not show_progress + disable=not show_progress, ): Y[i - self.step_interval + 1:i + 1] = self.step( T_step=self.T[i - self.step_interval:i + 1], y_0=Y[i - self.step_interval], - params=params + params=params, )[1:] return Y diff --git a/quantrl/solvers/context_manager.py b/quantrl/solvers/context_manager.py index 7259b97..4dfbaa0 100644 --- a/quantrl/solvers/context_manager.py +++ b/quantrl/solvers/context_manager.py @@ -6,28 +6,33 @@ __name__ = 'quantrl.solvers.context_manager' __authors__ = ["Sampreet Kalita"] __created__ = "2024-10-09" -__updated__ = "2025-05-11" +__updated__ = "2025-05-29" + +# dependencies +from typing import Type # quantrl modules from .base import BaseIVPSolver SOLVERS_IVP = {} -def get_IVP_solver( +def get_solver_ivp( library:str - ) -> BaseIVPSolver: + ) -> Type[BaseIVPSolver]: """Method to obtain an IVP solver class. Parameters ---------- library: str - Name of the library. Options are ``'jax'``, ``'torch'`` and ``'numpy'``. + Name of the library for the solver. + Options are ``'jax'``, ``'torch'`` and ``'numpy'``. Returns ------- IVPSolver: :class:`quantrl.solvers.base.BaseIVPSolver` The IVP solver class. """ + if library in SOLVERS_IVP: return SOLVERS_IVP[library] @@ -48,7 +53,8 @@ def get_IVP_solver( library = 'torch' return SOLVERS_IVP[library] - assert 'numpy' in library.lower(), "parameter ``library`` can be either ``'jax'`, ``'torch'`` or ``'numpy'``" + assert 'numpy' in library.lower(), \ + "parameter ``library`` can be either ``'jax'``, ``'torch'`` or ``'numpy'``" from .numpy import SciPyIVPSolver SOLVERS_IVP['numpy'] = SciPyIVPSolver library = 'numpy' diff --git a/quantrl/solvers/jax.py b/quantrl/solvers/jax.py index 8da43b6..59e2ba5 100644 --- a/quantrl/solvers/jax.py +++ b/quantrl/solvers/jax.py @@ -6,7 +6,7 @@ __name__ = 'quantrl.solvers.jax' __authors__ = ["Sampreet Kalita"] __created__ = "2024-03-10" -__updated__ = "2025-05-11" +__updated__ = "2025-05-29" # dependencies import jax @@ -19,26 +19,34 @@ # TODO: Implement interpolation class DiffraxIVPSolver(BaseIVPSolver): - """ODE and DDE solver using Diffrax-based methods for initial-value problems. + """ODE and DDE solver using + Diffrax-based methods + for initial-value problems. Available methods are ``'dopri5'``, ``'dopri8'``, and ``'tsit5'``. - Refer to :class:`quantrl.backends.base.BaseIVPSolver` for its implementation. + Refer to :class:`quantrl.backends.base.BaseIVPSolver` + for its implementation. """ # attributes - solver_methods = ['dopri5', 'dopri8', 'tsit5'] + solver_methods = [ + 'dopri5', + 'dopri8', + 'tsit5', + ] """list: Diffrax-based methods.""" - def __init__(self, - func, - y_0, - T, - solver_params:dict, - func_controls=None, - has_delay:bool=False, - func_delay=None, - delay_interval:int=0, - backend:JAXBackend=None + def __init__( + self, + func, + y_0, + T, + solver_params:dict, + func_controls=None, + has_delay:bool=False, + func_delay=None, + delay_interval:int=0, + backend:JAXBackend=None, ): # initialize BaseIVPSolver super().__init__( @@ -51,8 +59,8 @@ def __init__(self, func_delay=jax.jit(func_delay) if func_delay is not None else None, delay_interval=delay_interval, backend=backend if backend is not None else JAXBackend( - precision='double' - ) + precision='double', + ), ) # initialize solver @@ -60,17 +68,18 @@ def __init__(self, self.solver = { 'dopri5': dfx.Dopri5, 'dopri8': dfx.Dopri8, - 'tsit5': dfx.Tsit5 + 'tsit5': dfx.Tsit5, }.get(self.solver_params['method'], dfx.Dopri5)() - def integrate(self, - T_step, - y_0, - params=None + def integrate( + self, + T_step, + y_0, + params=None, ): # convert to tensor y_0 = self.backend.convert_to_typed( - tensor=y_0 + tensor=y_0, ) # integrate @@ -85,12 +94,13 @@ def integrate(self, saveat=dfx.SaveAt(ts=T_step), stepsize_controller=dfx.PIDController( atol=self.solver_params['atol'], - rtol=self.solver_params['rtol'] - ) + rtol=self.solver_params['rtol'], + ), ).ys - def interpolate(self, - T_step, - Y + def interpolate( + self, + T_step, + Y, ): raise NotImplementedError diff --git a/quantrl/solvers/numpy.py b/quantrl/solvers/numpy.py index 9e74b54..cf1a6ef 100644 --- a/quantrl/solvers/numpy.py +++ b/quantrl/solvers/numpy.py @@ -6,7 +6,7 @@ __name__ = 'quantrl.solvers.numpy' __authors__ = ["Sampreet Kalita"] __created__ = "2024-03-10" -__updated__ = "2025-05-11" +__updated__ = "2025-05-29" # dependencies import scipy.integrate as si @@ -17,30 +17,49 @@ from .base import BaseIVPSolver class SciPyIVPSolver(BaseIVPSolver): - """ODE and DDE solver using SciPy-based methods for initial-value problems. - - Available methods are ``'BDF'``, ``'DOP853'``, ``'LSODA'``, ``'Radau'``, ``'RK23'``, ``'RK45'``, ``'dop853'``, ``'dopri5'``, ``'lsoda'``, ``'vode'`` and ``'zvode'``. - Refer to :class:`quantrl.backends.base.BaseIVPSolver` for its implementation. + """ODE and DDE solver using + SciPy-based methods + for initial-value problems. + + Available methods are ``'BDF'``, ``'DOP853'``, ``'LSODA'``, + ``'Radau'``, ``'RK23'``, ``'RK45'``, ``'dop853'``, + ``'dopri5'``, ``'lsoda'``, ``'vode'`` and ``'zvode'``. + Refer to :class:`quantrl.backends.base.BaseIVPSolver` + for its implementation. """ # attributes - scipy_new_methods = ['BDF', 'DOP853', 'LSODA', 'Radau', 'RK23', 'RK45'] - """list: New Python-based methods availabile in :class:`scipy.integrate`.""" - scipy_old_methods = ['dop853', 'dopri5', 'lsoda', 'vode', 'zvode'] - """list: Old FORTRAN-based methods availabile in :class:`scipy.integrate`.""" + scipy_new_methods = [ + 'BDF', + 'DOP853', + 'LSODA', + 'Radau', + 'RK23', + 'RK45', + ] + """list: Python-based methods availabile in :class:`scipy.integrate`.""" + scipy_old_methods = [ + 'dop853', + 'dopri5', + 'lsoda', + 'vode', + 'zvode', + ] + """list: FORTRAN-based methods availabile in :class:`scipy.integrate`.""" solver_methods = scipy_new_methods + scipy_old_methods """list: SciPy-based methods availabile in :class:`scipy.integrate`.""" - def __init__(self, - func, - y_0, - T, - solver_params:dict, - func_controls=None, - has_delay:bool=False, - func_delay=None, - delay_interval:int=0, - backend:NumPyBackend=None + def __init__( + self, + func, + y_0, + T, + solver_params:dict, + func_controls=None, + has_delay:bool=False, + func_delay=None, + delay_interval:int=0, + backend:NumPyBackend=None, ): # initialize BaseIVPSolver super().__init__( @@ -53,8 +72,8 @@ def __init__(self, func_delay=func_delay, delay_interval=delay_interval, backend=backend if backend is not None else NumPyBackend( - precision='double' - ) + precision='double', + ), ) # flatten function for integration @@ -68,8 +87,8 @@ def __init__(self, tensor=y, shape=self.shape_y ), - args=args - ) + args=args, + ), ) self.is_y_flat = False @@ -81,43 +100,45 @@ def __init__(self, name=self.solver_params['method'], atol=self.solver_params['atol'], rtol=self.solver_params['rtol'], - method='bdf' if self.solver_params['is_stiff'] else 'adams' + method='bdf' if self.solver_params['is_stiff'] else 'adams', ) - def integrate(self, - T_step, - y_0, - params=None + def integrate( + self, + T_step, + y_0, + params=None, ): # convert to tensor y_0 = self.backend.convert_to_typed( - tensor=y_0 + tensor=y_0, ) # flatten y_0_flat = y_0 if not self.is_y_flat: y_0_flat = self.backend.flatten( - tensor=y_0 + tensor=y_0, ) # integrate _Y_flat = self.integrate_flat( y_0_flat=y_0_flat, T_step=T_step, - args=[params, self.func_controls, self.func_delay] + args=[params, self.func_controls, self.func_delay], ) # reshape return self.backend.reshape( tensor=_Y_flat, - shape=(len(T_step), *self.shape_y) + shape=(len(T_step), *self.shape_y), ) - def integrate_flat(self, - T_step, - y_0_flat, - args:tuple + def integrate_flat( + self, + T_step, + y_0_flat, + args:tuple, ): """Method to take one integration step. @@ -135,6 +156,7 @@ def integrate_flat(self, Y: Any Values of the variables at the given points of time. """ + # convert to tensor y_0_flat = self.backend.convert_to_typed( tensor=y_0_flat @@ -147,14 +169,14 @@ def integrate_flat(self, *self.backend.shape( tensor=T_step ), - *y_0_flat.shape + *y_0_flat.shape, ), - dtype='complex' if self.solver_params['complex'] else 'real' + dtype='complex' if self.solver_params['complex'] else 'real', ) _Y_flat[0] = y_0_flat self.integrator.set_initial_value( y=y_0_flat, - t=T_step[0] + t=T_step[0], ) self.integrator.set_f_params(args) for i in range(1, len(T_step)): @@ -171,18 +193,19 @@ def integrate_flat(self, method=self.solver_params['method'], atol=self.solver_params['atol'], rtol=self.solver_params['rtol'], - args=(args, ) - ).y + args=(args, ), + ).y, ) return _Y_flat - def interpolate(self, - T_step, - Y + def interpolate( + self, + T_step, + Y, ): _shape = self.backend.shape( - tensor=Y + tensor=Y, )[1] b_spline = [splrep(T_step, Y[:, j]) for j in range(_shape)] return lambda t: [splev(t, b_spline[j]) for j in range(_shape)] diff --git a/quantrl/solvers/torch.py b/quantrl/solvers/torch.py index 8fd86d0..57697be 100644 --- a/quantrl/solvers/torch.py +++ b/quantrl/solvers/torch.py @@ -6,7 +6,7 @@ __name__ = 'quantrl.solvers.torch' __authors__ = ["Sampreet Kalita"] __created__ = "2024-03-10" -__updated__ = "2025-05-11" +__updated__ = "2025-05-29" # dependencies from torchdiffeq import odeint @@ -18,26 +18,38 @@ # TODO: Implement interpolation class TorchDiffEqIVPSolver(BaseIVPSolver): - """ODE and DDE solver using TorchDiffEq-based methods for initial-value problems. + """ODE and DDE solver using + TorchDiffEq-based methods + for initial-value problems. - Available methods are ``'adaptive_huen'``, ``'bosh3'``, ``'dopri5'``, ``'dopri8'````'fehlberg2'`` and ``'tsit5'``. - Refer to :class:`quantrl.backends.base.BaseIVPSolver` for its implementation. + Available methods are ``'adaptive_huen'``, ``'bosh3'``, + ``'dopri5'``, ``'dopri8'````'fehlberg2'`` and ``'tsit5'``. + Refer to :class:`quantrl.backends.base.BaseIVPSolver` + for its implementation. """ # attributes - solver_methods = ['adaptive_huen', 'bosh3', 'dopri5', 'dopri8', 'fehlberg2', 'tsit5'] + solver_methods = [ + 'adaptive_huen', + 'bosh3', + 'dopri5', + 'dopri8', + 'fehlberg2', + 'tsit5', + ] """list: TorchDiffEq-based methods.""" - def __init__(self, - func, - y_0, - T, - solver_params:dict, - func_controls=None, - has_delay:bool=False, - func_delay=None, - delay_interval:int=0, - backend:TorchBackend=None + def __init__( + self, + func, + y_0, + T, + solver_params:dict, + func_controls=None, + has_delay:bool=False, + func_delay=None, + delay_interval:int=0, + backend:TorchBackend=None, ): # initialize BaseIVPSolver super().__init__( @@ -51,33 +63,39 @@ def __init__(self, delay_interval=delay_interval, backend=backend if backend is not None else TorchBackend( precision='double', - device='cuda' - ) + device='cuda', + ), ) - def integrate(self, - T_step, - y_0, - params=None + def integrate( + self, + T_step, + y_0, + params=None, ): # convert to tensor y_0 = self.backend.convert_to_typed( - tensor=y_0 + tensor=y_0, ) # integrate return odeint( - func=lambda t, y: self.func(t, y, [params, self.func_controls, self.func_delay]), + func=lambda t, y: self.func( + t, + y, + [params, self.func_controls, self.func_delay], + ), y0=y_0, t=T_step, atol=self.solver_params['atol'], rtol=self.solver_params['rtol'], method=self.solver_params['method'], - options={} + options={}, ) - def interpolate(self, - T_step, - Y + def interpolate( + self, + T_step, + Y, ): raise NotImplementedError diff --git a/quantrl/utils.py b/quantrl/utils.py index f96334f..c526daa 100644 --- a/quantrl/utils.py +++ b/quantrl/utils.py @@ -6,7 +6,7 @@ __name__ = 'quantrl.utils' __authors__ = ["Sampreet Kalita"] __created__ = "2023-06-02" -__updated__ = "2024-10-14" +__updated__ = "2025-08-20" # dependencies import gc @@ -37,13 +37,14 @@ class SaveOnBestMeanRewardCallback(BaseCallback): Number of episodes over which the average reward is calculated. """ - def __init__(self, - update_steps:int, - log_dir:str, - n_episodes:int, - episode_start:int, - steps_per_episode:int, - average_over:int=100 + def __init__( + self, + update_steps:int, + log_dir:str, + n_episodes:int, + episode_start:int, + steps_per_episode:int, + average_over:int=100, ): """Class constructor for SaveOnBestMeanRewardCallback.""" @@ -62,10 +63,18 @@ def __init__(self, self.t_start = time.time() # initialize file to save mean rewards - with open(self.log_dir + f'reward_mean_{self.episode_start}_{self.n_episodes - 1}.txt', 'w', encoding='utf-8') as file: - s = f'{"time":>14} {"n_steps":>12} {"episode_curr":>12} {"reward_curr":>14} {"episode_best":>12} {"reward_best":>14}\n' - file.write(s) - file.close() + with open( + self.log_dir \ + + f"reward_mean_{self.episode_start}_{self.n_episodes - 1}" \ + + ".txt", + 'w', + encoding='utf-8', + ) as file_handler: + s = f"{"time":>14} {"n_steps":>12} {"episode_curr":>12} " \ + + f"{"reward_curr":>14} {"episode_best":>12} " \ + + f"{"reward_best":>14}\n" + file_handler.write(s) + file_handler.close() def _init_callback(self) -> None: """Method on intialization.""" @@ -80,9 +89,19 @@ def _on_step(self) -> bool: if self.n_calls % self.update_steps == 0: # retrieve reward data from monitor file - with open(self.log_dir + f'learning_{self.episode_start}_{self.n_episodes - 1}.monitor.csv', 'r', encoding='utf-8') as file_handler: + with open( + self.log_dir \ + + f"learning_{self.episode_start}_{self.n_episodes - 1}" \ + + ".monitor.csv", + 'r', + encoding='utf-8', + ) as file_handler: file_handler.readline() - xs, ys = ts2xy(pandas.read_csv(file_handler, index_col=None), 'timesteps') + xs, ys = ts2xy( + pandas.read_csv(file_handler, index_col=None), + 'timesteps', + ) + file_handler.close() if len(xs) > 0: # current mean reward @@ -92,7 +111,9 @@ def _on_step(self) -> bool: # update console if self.verbose >= 1: - print(f"Best mean reward: {self.reward_best:.6f} at #{self.episode_best:6d}\nCurr mean reward: {reward_curr:.6f}") + print(f"Best mean reward: {self.reward_best:.6f} " \ + + f"at #{self.episode_best:6d}\n" \ + + f"Curr mean reward: {reward_curr:.6f}") # New best model, you could save the agent here if reward_curr > self.reward_best: @@ -104,12 +125,25 @@ def _on_step(self) -> bool: print("Saving new best model and replay buffer...") # save model and replay buffer - self.model.save(self.log_dir + f'models/best_{self.n_episodes - 1}.zip') - self.model.save_replay_buffer(self.log_dir + f'buffers/best_{self.n_episodes - 1}.zip') + self.model.save(self.log_dir \ + + f"models/best_{self.n_episodes - 1}.zip") + self.model.save_replay_buffer(self.log_dir \ + + f"buffers/best_{self.n_episodes - 1}.zip") # save reward data - with open(self.log_dir + f'reward_mean_{self.episode_start}_{self.n_episodes - 1}.txt', 'a', encoding='utf-8') as file: - file.write(f'{time.time() - self.t_start:14.03f} {self.n_calls:12d} {episode_curr:12d} {reward_curr:14.06f} {self.episode_best:12d} {self.reward_best:14.06f}\n') - file.close() + with open( + self.log_dir \ + + f"reward_mean_{self.episode_start}_" \ + + f"{self.n_episodes - 1}.txt", + 'a', + encoding='utf-8', + ) as file_handler: + file_handler.write(f"{time.time() - self.t_start:14.03f}" \ + + f"{self.n_calls:12d}" \ + + f"{episode_curr:12d}" \ + + f"{reward_curr:14.06f}" \ + + f"{self.episode_best:12d}" \ + + f"{self.reward_best:14.06f}\n") + file_handler.close() gc.collect() return True From aa0b3390c7c419eaf56d43c7ab4b859f3cae4f85 Mon Sep 17 00:00:00 2001 From: Sampreet Kalita <9553215+Sampreet@users.noreply.github.com> Date: Sun, 26 Jul 2026 22:36:49 +0100 Subject: [PATCH 4/5] Add Support for Tuple Spaces --- CHANGELOG.md | 5 ++ quantrl/backends/torch.py | 4 +- quantrl/envs/base.py | 89 +++++++++++++++++++++++++----- quantrl/solvers/base.py | 9 ++- quantrl/solvers/context_manager.py | 2 +- quantrl/solvers/jax.py | 3 +- quantrl/utils.py | 5 +- 7 files changed, 96 insertions(+), 21 deletions(-) diff --git a/CHANGELOG.md b/CHANGELOG.md index d3efa1c..57cea1b 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -1,5 +1,10 @@ # Changelog +## 2026/07/16 - 00 - v0.0.11 - Add Support for Tuple Spaces +* Added support for tuple spaces to combine discrete and continuious actions in `quantrl.envs.base` module. +* Added option for maximum number of steps for `DiffraxIVPSolver` in `quantrl.solver.jax` module. +* Minor updates to `quantrl.backends.torch` module. + ## 2025/08/20 - 00 - v0.0.10 - Code Cleanup * Minor changes to all modules except `quantrl.solvers.measure`. * Updated indentations and line lengths of modules and `README`. diff --git a/quantrl/backends/torch.py b/quantrl/backends/torch.py index a18abae..ea57931 100644 --- a/quantrl/backends/torch.py +++ b/quantrl/backends/torch.py @@ -71,12 +71,12 @@ def convert_to_typed( def convert_to_numpy( self, - tensor, + tensor:torch.Tensor, dtype:str=None, ) -> np.ndarray: if self.is_typed( tensor=tensor, - dtype=dtype, + dtype=None, ): return np.asarray( tensor.detach().cpu().numpy() \ diff --git a/quantrl/envs/base.py b/quantrl/envs/base.py index 4e7963e..bbf3a94 100644 --- a/quantrl/envs/base.py +++ b/quantrl/envs/base.py @@ -6,14 +6,14 @@ __name__ = 'quantrl.envs.base' __authors__ = ["Sampreet Kalita"] __created__ = "2023-04-25" -__updated__ = "2025-08-20" +__updated__ = "2026-07-16" # dependencies from abc import ABC, abstractmethod import sys from gymnasium import Env -from gymnasium.spaces import Box, MultiDiscrete +from gymnasium.spaces import Box, MultiDiscrete, Tuple import numpy as np from stable_baselines3.common import env_util from stable_baselines3.common.vec_env import VecEnv @@ -86,8 +86,17 @@ class BaseEnv(ABC): is scaled by the corresponding action multiplier. Default is ``[-1.0, 1.0]``. action_space_type (*str*) the type of action space. - Options are ``'binary'`` and ``'box'``. + Options are ``'discrete'``, ``'multidiscrete'``, + ``'box'`` and ``'tuple'``. Default is ``'box'``. + action_space_seed (*int*) the seed for the action space. + Default is ``None``. + action_space_tuple (*tuple*) dictionaries of action spaces containing + the keys ``'type'`` for the type of action space and + ``'range'`` for the range of the action space. + Currently, only ``'discrete'`` and ``'box'`` types + are supported for tuples. Default is a dictionary + with the parameters of a box action space. seed (*int*) seed to initialize random number generators. If ``None``, a random integer seed is generated. Default is ``None``. @@ -142,6 +151,12 @@ class BaseEnv(ABC): 'observation_stds': None, 'action_space_range': [-1.0, 1.0], 'action_space_type': 'box', + 'action_space_seed': None, + 'action_space_tuple': ({ + 'type': 'box', + 'n_actions': 1, + 'range': [-1.0, 1.0], + }, ), 'seed': None, 'cache_all_data': True, 'cache_dump_interval': 100, @@ -192,8 +207,8 @@ def __init__( "parameter ``n_properties`` should be non-negative" assert action_interval > 0, \ "parameter ``action_interval`` should be a positive integer" - assert len(data_idxs) > 0, \ - "parameter ``data_idxs`` should be a list containing at least one element" + assert data_idxs is None or len(data_idxs) > 0, \ + "parameter ``data_idxs`` should be a list containing at least one element or ``None``" assert kwargs['observation_stds'] is None \ or isinstance(kwargs['observation_stds'], list), \ "parameter ``observation_stds`` should be a list" @@ -213,8 +228,9 @@ def __init__( "parameter ``action_space_range`` should contain " \ + "two elements for the minimum and maximum values, " \ + "both inclusive" - assert kwargs['action_space_type'] in ['binary', 'box'], \ - "parameter ``action_space_type`` can be either ``'binary'`` or ``'box'``" + assert kwargs['action_space_type'] in ['discrete', 'multidiscrete', 'box', 'tuple'], \ + "parameter ``action_space_type`` can be either \ + ``'discrete'``, ``'multidiscrete'``, ``'box'`` or ``'tuple'``" assert kwargs['cache_dump_interval'] > 0, \ "parameter ``cache_dump_interval`` should be a positive integer" @@ -264,24 +280,70 @@ def __init__( # action attributes self.n_actions = n_actions self.action_space_type = kwargs['action_space_type'] + self.action_space_seed = kwargs['action_space_seed'] # discrete actions - if self.action_space_type == 'binary': - self.action_space_range = [0, 1] + if self.action_space_type == 'discrete': + _range = kwargs['action_space_range'] + _diff = _range[1] - _range[0] self.action_space = MultiDiscrete( - nvec=[2] * self.n_actions, + nvec=[_diff] * self.n_actions, + dtype=self.numpy_int, + seed=self.action_space_seed, + start=[_range[0]] * self.n_actions, + ) + elif self.action_space_type == 'multidiscrete': + _ranges = kwargs['action_space_ranges'] + _diffs = [_range[1] - _range[0] for _range in _ranges] + self.action_space = MultiDiscrete( + nvec=_diffs, + dtype=self.numpy_int, + seed=self.action_space_seed, + start=[_range[0] for _range in _ranges], ) # continuous actions - else: + elif self.action_space_type == 'box': self.action_space_range = kwargs['action_space_range'] self.action_space = Box( low=self.action_space_range[0], high=self.action_space_range[1], shape=(self.n_actions, ), dtype=self.numpy_real, + seed=self.action_space_seed, + ) + # custom ordered actions + else: + self.action_space_tuple = kwargs['action_space_tuple'] + spaces = [] + _n_actions = 0 + for entry in self.action_space_tuple: + if entry['type'] == 'discrete': + _diff = entry['range'][1] - entry['range'][0] + space = MultiDiscrete( + nvec=[_diff] * entry['n_actions'], + dtype=self.numpy_int, + seed=self.action_space_seed, + start=[entry['range'][0]] * entry['n_actions'], + ) + else: + space = Box( + low=entry['range'][0], + high=entry['range'][1], + shape=(entry['n_actions'], ), + dtype=self.numpy_real, + seed=self.action_space_seed, + ) + _n_actions += entry['n_actions'] + spaces.append(space) + assert _n_actions == self.n_actions, \ + "sum of the number of actions of all subspaces \ + should be equal to the total number of actoins ``n_actions``" + self.action_space = Tuple( + spaces=spaces, + seed=self.action_space_seed, ) self.action_maximums = self.backend.convert_to_typed( tensor=action_maximums, - dtype='integer' if self.action_space_type == 'binary' else 'real' + dtype='integer' if self.action_space_type == 'discrete' else 'real' ) self.action_interval = action_interval self.action_steps = self.numpy_int( @@ -312,7 +374,8 @@ def __init__( self.average_over = self.numpy_int(kwargs['average_over']) # initialize IO - self.data_idxs = data_idxs + self.data_idxs = data_idxs if data_idxs is not None else \ + list(range(self.n_data)) self.cache_all_data = kwargs['cache_all_data'] self.cache_dump_interval = kwargs['cache_dump_interval'] self.io = FileIO( diff --git a/quantrl/solvers/base.py b/quantrl/solvers/base.py index f8007be..1b71f83 100644 --- a/quantrl/solvers/base.py +++ b/quantrl/solvers/base.py @@ -6,7 +6,7 @@ __name__ = 'quantrl.solvers.base' __authors__ = ["Sampreet Kalita"] __created__ = "2024-03-10" -__updated__ = "2025-08-20" +__updated__ = "2026-07-16" # dependencies from abc import ABC, abstractmethod @@ -54,6 +54,10 @@ class BaseIVPSolver(ABC): step_interval (*bool*) number of steps to jump during integration. Higher values give faster results. Default is ``10``. + step_max (*int*) maximum number of steps to take + before the computation is terminated, + used by :class:`quantrl.solvers.jax.DiffraxIVPSolver`. + Default is ``10_000``. ================ ============================================ func_controls: callable Function for the controls in the format ``func_controls(t)``. @@ -79,6 +83,7 @@ class BaseIVPSolver(ABC): 'is_stiff': False, 'step_interval': 10, 'complex': False, + 'step_max': 10_000, } """dict: Default parameters of the solver.""" solver_methods = [] @@ -127,7 +132,7 @@ def __init__( for key, _ in self.default_solver_params.items(): self.solver_params[key] = solver_params.get( key, - self.default_solver_params[key] + self.default_solver_params[key], ) # override step dimension with delay interval if DDE if self.has_delay and self.delay_interval != 0: diff --git a/quantrl/solvers/context_manager.py b/quantrl/solvers/context_manager.py index 4dfbaa0..3ecfb11 100644 --- a/quantrl/solvers/context_manager.py +++ b/quantrl/solvers/context_manager.py @@ -44,7 +44,7 @@ def get_solver_ivp( return SOLVERS_IVP[library] # use PyTorch if JAX is not installed except ImportError: - print("JAX not installed, defaulting to PyTorch") + print("JAX or Diffrax not installed, defaulting to PyTorch") library = 'torch' if 'torch' in library.lower(): diff --git a/quantrl/solvers/jax.py b/quantrl/solvers/jax.py index 59e2ba5..4463a76 100644 --- a/quantrl/solvers/jax.py +++ b/quantrl/solvers/jax.py @@ -6,7 +6,7 @@ __name__ = 'quantrl.solvers.jax' __authors__ = ["Sampreet Kalita"] __created__ = "2024-03-10" -__updated__ = "2025-05-29" +__updated__ = "2026-07-17" # dependencies import jax @@ -96,6 +96,7 @@ def integrate( atol=self.solver_params['atol'], rtol=self.solver_params['rtol'], ), + max_steps=self.solver_params['step_max'], ).ys def interpolate( diff --git a/quantrl/utils.py b/quantrl/utils.py index c526daa..e111af4 100644 --- a/quantrl/utils.py +++ b/quantrl/utils.py @@ -127,8 +127,9 @@ def _on_step(self) -> bool: # save model and replay buffer self.model.save(self.log_dir \ + f"models/best_{self.n_episodes - 1}.zip") - self.model.save_replay_buffer(self.log_dir \ - + f"buffers/best_{self.n_episodes - 1}.zip") + if getattr(self.model, "save_replay_buffer", None) is not None: + self.model.save_replay_buffer(self.log_dir \ + + f"buffers/best_{self.n_episodes - 1}.zip") # save reward data with open( From f1a587708a6a342f8ccd0183f9ef2b5ac7adc4a5 Mon Sep 17 00:00:00 2001 From: Sampreet Kalita <9553215+Sampreet@users.noreply.github.com> Date: Wed, 2 Sep 2026 15:20:31 +0100 Subject: [PATCH 5/5] Bump version and Minor Fixes --- CHANGELOG.md | 7 +++++++ README.md | 12 ++++++------ docs/source/conf.py | 2 +- quantrl/__init__.py | 2 +- quantrl/backends/jax.py | 4 ++-- quantrl/envs/base.py | 24 ++++++++++++++---------- quantrl/envs/stochastic.py | 4 ++-- quantrl/plotters.py | 19 ++++++++++++++----- 8 files changed, 47 insertions(+), 27 deletions(-) diff --git a/CHANGELOG.md b/CHANGELOG.md index 57cea1b..a014a0a 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -1,5 +1,12 @@ # Changelog +## 2026/09/02 - 00 - v0.1.0 - Bump version and Minor Fixes +* Minor fixes to `quantrl.backends.jax`, `quantrl.envs.base` and `quantrl.envs.stochastic` modules. +* Renamed ``seed`` to ``noise_seed`` in `quantrl.envs.base` module. +* Updated ``LearningCurvePlotter`` class in `quantrl.plotters` module. +* Bump version in `docs/source/conf.py`, `quantrl.__init__.py`. +* Updated `README`. + ## 2026/07/16 - 00 - v0.0.11 - Add Support for Tuple Spaces * Added support for tuple spaces to combine discrete and continuious actions in `quantrl.envs.base` module. * Added option for maximum number of steps for `DiffraxIVPSolver` in `quantrl.solver.jax` module. diff --git a/README.md b/README.md index ea25c37..11ebacf 100644 --- a/README.md +++ b/README.md @@ -1,6 +1,6 @@ # QuantRL: Quantum Control using Reinforcement Learning -![Latest Version](https://img.shields.io/badge/version-0.0.10-red?style=for-the-badge) +![Latest Version](https://img.shields.io/badge/version-0.1.0-red?style=for-the-badge) > A backend-agnostic library of modules to interface deterministic and stochastic quantum models for reinforcement learning. @@ -17,8 +17,8 @@ with asynchronous saves. ### What's New! -* Support for NumPy 2.x.x. -* ``'tsit5'`` solver in PyTorch. +* Support for PyTorch (GPU) with ``'tsit'`` solver. +* Support for tuple action spaces. For a complete list of changes, see [CHANGELOG.md](CHANGELOG.md). @@ -26,7 +26,7 @@ For a complete list of changes, see [CHANGELOG.md](CHANGELOG.md). [QuantRL](https://github.com/sampreet/quantrl) requires `Python 3.12+`, preferably installed via the -[Anaconda distribution](https://www.anaconda.com/download). +[MiniForge distribution](https://conda-forge.org/download/). It's base dependencies can be installed using: ```bash @@ -49,11 +49,11 @@ For the CPU versions, use: python -m pip install torch torchdiffeq jax diffrax ``` -For the GPU versions with CUDA 12 support, use: +For the GPU versions with CUDA 13 support, use: ```bash python -m pip install torch --index-url https://download.pytorch.org/whl/cu126 -python -m pip install torchdiffeq "jax[cuda12]" diffrax +python -m pip install torchdiffeq "jax[cuda13]" diffrax ``` ***Note: JAX-GPU support for Windows and MacOS diff --git a/docs/source/conf.py b/docs/source/conf.py index 5379f07..ba61672 100644 --- a/docs/source/conf.py +++ b/docs/source/conf.py @@ -22,7 +22,7 @@ author = 'Sampreet Kalita' # The full version, including alpha/beta/rc tags -release = '0.0.10' +release = '0.1.0' # -- General configuration --------------------------------------------------- diff --git a/quantrl/__init__.py b/quantrl/__init__.py index 017dc0d..912bb12 100644 --- a/quantrl/__init__.py +++ b/quantrl/__init__.py @@ -1,2 +1,2 @@ """Module to initialize QuantRL.""" -__version__ = "0.0.10" +__version__ = "0.1.0" diff --git a/quantrl/backends/jax.py b/quantrl/backends/jax.py index c74cdfe..21cbdc3 100644 --- a/quantrl/backends/jax.py +++ b/quantrl/backends/jax.py @@ -6,7 +6,7 @@ __name__ = 'quantrl.backends.jax' __authors__ = ["Sampreet Kalita"] __created__ = "2024-03-10" -__updated__ = "2025-06-04" +__updated__ = "2026-08-11" # dependencies from inspect import getfullargspec @@ -65,7 +65,7 @@ def transpose( # get swapped axes _shape = jnp.shape(tensor) - _axes = jnp.arange(len(_shape)) + _axes = list(range(len(_shape))) _axes[axis_1] = axis_0 % len(_shape) _axes[axis_0] = axis_1 % len(_shape) diff --git a/quantrl/envs/base.py b/quantrl/envs/base.py index bbf3a94..3a33b2f 100644 --- a/quantrl/envs/base.py +++ b/quantrl/envs/base.py @@ -6,7 +6,7 @@ __name__ = 'quantrl.envs.base' __authors__ = ["Sampreet Kalita"] __created__ = "2023-04-25" -__updated__ = "2026-07-16" +__updated__ = "2026-08-13" # dependencies from abc import ABC, abstractmethod @@ -97,7 +97,7 @@ class BaseEnv(ABC): Currently, only ``'discrete'`` and ``'box'`` types are supported for tuples. Default is a dictionary with the parameters of a box action space. - seed (*int*) seed to initialize random number + noise_seed (*int*) seed to initialize random number generators. If ``None``, a random integer seed is generated. Default is ``None``. cache_all_data (*bool*) option to cache all data to disk. @@ -355,7 +355,7 @@ def __init__( self.t_delay = self.T[self.action_interval] - self.T[0] # initialize seed - self.seed = kwargs['seed'] + self.noise_seed = kwargs['seed'] # data constants self.dir_path = dir_prefix + "/" + "_".join([ @@ -422,7 +422,7 @@ def _update_states(self): States: Any The updated states with shape either ``(action_interval + 1, n_observations)`` or - ``(action_interval + 1, n_envs, n_observations). + ``(action_interval + 1, n_envs, n_observations)``. """ raise NotImplementedError @@ -571,7 +571,7 @@ def reset(self): if self.observation_stds is not None: self.Observation_noises = self.backend.normal( generator=self.backend.generator( - seed=self.seed, + seed=self.noise_seed, ), shape=(self.shape_T[0], *_shape), mean=0.0, @@ -722,14 +722,14 @@ def plot_learning_curve( # get reward data from file if data_rewards is None: data_rewards = self.io.load_data( - file_name=file_name + file_name=file_name, ) # get reward data from trajectories if data_rewards is None: data_rewards = self.io.get_disk_cache( idx_start=_idx_s, idx_end=_idx_e, - idxs=[-1] + idxs=[-1], )[:, -1, 0] # initialize plotter @@ -1029,7 +1029,7 @@ def step( # check if truncation required truncated = self.check_truncation() if truncated > 0: - print(f"Trajectory #{self.traj_idx} truncated") + print(f"Trajectory #{self.traj_idx} truncated after {self.t_idx} steps") # if trajectory ends if terminated or truncated: @@ -1564,8 +1564,12 @@ def step_wait(self): # reset variables observations = self._reset() - return observations, reward, \ - [terminated or truncated] * self.n_envs, [{}] * self.n_envs + dones = self.backend.convert_to_numpy( + tensor=[terminated or truncated] * self.n_envs, + dtype='integer', + ) + + return observations, reward, dones, [{}] * self.n_envs def update_data(self): """Method to update the batch data for the step. diff --git a/quantrl/envs/stochastic.py b/quantrl/envs/stochastic.py index b31c312..d0669e7 100644 --- a/quantrl/envs/stochastic.py +++ b/quantrl/envs/stochastic.py @@ -6,7 +6,7 @@ __name__ = 'quantrl.envs.stochastic' __authors__ = ["Sampreet Kalita"] __created__ = "2023-04-25" -__updated__ = "2025-08-20" +__updated__ = "2026-08-11" # dependencies import numpy as np @@ -185,7 +185,7 @@ def reset( # update Wiener noises self.Ws = np.sqrt(self.t_ssz) * self.backend.normal( generator=self.backend.generator( - seed=self.seed, + seed=self.noise_seed, ), shape=(self.shape_T[0] - 1, self.n_observations), mean=0.0, diff --git a/quantrl/plotters.py b/quantrl/plotters.py index d49495a..589ad3e 100644 --- a/quantrl/plotters.py +++ b/quantrl/plotters.py @@ -6,7 +6,7 @@ __name__ = 'quantrl.plotters' __authors__ = ["Sampreet Kalita"] __created__ = "2023-12-08" -__updated__ = "2025-08-20" +__updated__ = "2026-09-02" # dependencies from io import BytesIO @@ -271,6 +271,9 @@ class LearningCurvePlotter(): percentiles: list, default=None Percentile values for intraquartile ranges. If ``None``, the percentiles are set to ``[25, 50, 75]``. + kwargs : dict + Other keyword arguments for figure and label and tick font properties + (with keys ``"label_fontdict"`` and ``"tick_fontdict"``). """ def __init__( @@ -278,6 +281,7 @@ def __init__( axis_args:list, average_over:int=100, percentiles:list=None, + **kwargs, ): """Class constructor for LearningCurvePlotter.""" @@ -292,12 +296,17 @@ def __init__( plt.ion() # initialize plot - self.fig = plt.figure(figsize=(6.0, 3.0)) + figsize = kwargs.get("figsize", (6.0, 3.0)) + label_fontdict = kwargs.get("label_fontdict", None) + tick_fontdict = kwargs.get("tick_fontdict", None) + self.fig = plt.figure(figsize=figsize) self.ax = plt.gca() - self.ax.set_xlabel(self.axis_args[0]) - self.ax.set_ylabel(self.axis_args[1]) + self.ax.set_xlabel(self.axis_args[0], fontdict=label_fontdict) + self.ax.set_ylabel(self.axis_args[1], fontdict=label_fontdict) self.ax.set_ylim(ymin=self.axis_args[2][0], ymax=self.axis_args[2][1]) self.ax.set_yscale(self.axis_args[3]) + if isinstance(tick_fontdict, dict): + self.ax.tick_params(labelsize=tick_fontdict.get("size", 14)) self.fig.tight_layout() # initialze buffer @@ -376,7 +385,7 @@ def add_data( q_min, q_max, facecolor=color, - alpha=0.1, + alpha=0.2, ) # calculate mean else: