Merge software changes into paper branch (#6)

* Add thin and progress kwargs to fitters

* Create test_package workflow (#4)

* Create test_package workflow

* Skip unit tests requiring mpi4py if not installed

* Add manual testing instructions to README

* Create CONTRIBUTING.md and move testing instructions here

* Link to contributing guidelines in docs

Author: AlanPearl

.github/workflows/test_package.yml

@@ -0,0 +1,46 @@
+name: Test diffopt
+
+on:
+  workflow_dispatch: null
+  schedule:
+    # Runs "every Monday at noon UTC"
+    - cron: '0 12 * * 1'
+  push:
+    branches:
+      - main
+  pull_request: null
+
+jobs:
+  build:
+    runs-on: ubuntu-latest
+    strategy:
+      matrix:
+        python-version: ['3.9', '3.10', '3.11']
+    steps:
+    - uses: actions/checkout@v2
+    - name: Set up Python ${{ matrix.python-version }}
+      uses: conda-incubator/setup-miniconda@v2
+      with:
+        activate-environment: test${{ matrix.python-version}}
+        python-version: ${{ matrix.python-version }}
+        channels: conda-forge,defaults
+        channel-priority: strict
+        show-channel-urls: true
+        miniforge-version: latest
+
+    - name: Install dependencies
+      shell: bash -l {0}
+      run: |
+        conda install -yq jax 
+        conda install -yq pip pytest pytest-cov flake8
+
+    - name: Install package
+      shell: bash -l {0}
+      run: |
+        pip install -e .
+
+    - name: Run tests
+      shell: bash -l {0}
+      run: |
+        export PYTHONWARNINGS=error
+        pytest -v

CONTRIBUTING.md

@@ -0,0 +1,36 @@
+# Contributing to `diffopt`
+
+Thank you for your interest in contributing to this project. All questions and ideas for improvement are welcome and can be made through opening an issue or pull request.
+
+Before contributing, familiarize yourself with our resources:
+
+- [Source Code](https://github.com/AlanPearl/diffopt)
+- [Documentation](https://diffopt.readthedocs.io)
+
+## Issues
+
+You can open an [issue](https://github.com/AlanPearl/diffopt/issues) if you:
+
+- Have encountered a bug or issue when using the software
+- Would like to see a new feature
+- Are seeking support that could not be resolved by reading the documentation
+
+## Pull Requests
+
+If you would like to directly submit your own change to the software, thank you! Here's how:
+
+- Fork [this repository](https://github.com/AlanPearl/diffopt).
+- Please remember to include a concise, self-contained unit test in your pull request. Ensure that all tests pass (see [Manual Testing](#manual-testing)).
+- Open a [pull request](https://github.com/AlanPearl/diffopt/pulls).
+
+## Manual Testing
+
+Make sure you have installed diffopt as described in the [docs](https://diffopt.readthedocs.io/en/latest/installation.html). To run all tests from the main directory:
+
+```bash
+pip install pytest
+pytest .
+mpirun -n 2 pytest .
+```
+
+Note that unit tests requiring `mpi4py` installation are not automatically tested by GitHub workflows. Therefore, running these tests manually with `mpi4py` installed is necessary to assure that all tests pass.

diffopt/kdescent/descent.py

@@ -12,7 +12,8 @@
 
 
 def adam(lossfunc, guess, nsteps=100, param_bounds=None,
-         learning_rate=0.01, randkey=1, const_randkey=False, **other_kwargs):
+         learning_rate=0.01, randkey=1, const_randkey=False,
+         thin=1, progress=True, **other_kwargs):
     """
     Perform gradient descent
 
@@ -36,6 +37,11 @@ def adam(lossfunc, guess, nsteps=100, param_bounds=None,
     const_randkey : bool, optional
         By default (False), randkey is regenerated at each gradient descent
         iteration. Remove this behavior by setting const_randkey=True
+    thin : int, optional
+        Return parameters for every `thin` iterations, by default 1. Set
+        `thin=0` to only return final parameters
+    progress : bool, optional
+        Display tqdm progress bar, by default True
 
     Returns
     -------
@@ -46,7 +52,7 @@ def adam(lossfunc, guess, nsteps=100, param_bounds=None,
     if param_bounds is None:
         return adam_unbounded(
             lossfunc, guess, nsteps, learning_rate, randkey,
-            const_randkey, **other_kwargs)
+            const_randkey, thin, progress, **other_kwargs)
 
     assert len(guess) == len(param_bounds)
     if hasattr(param_bounds, "tolist"):
@@ -60,14 +66,15 @@ def ulossfunc(uparams, *args, **kwargs):
     init_uparams = apply_transforms(guess, param_bounds)
     uparams = adam_unbounded(
         ulossfunc, init_uparams, nsteps, learning_rate, randkey,
-        const_randkey, **other_kwargs)
+        const_randkey, thin, progress, **other_kwargs)
     params = apply_inverse_transforms(uparams.T, param_bounds).T
 
     return params
 
 
 def adam_unbounded(lossfunc, guess, nsteps=100, learning_rate=0.01,
-                   randkey=1, const_randkey=False, **other_kwargs):
+                   randkey=1, const_randkey=False,
+                   thin=1, progress=True, **other_kwargs):
     kwargs = {**other_kwargs}
     if randkey is not None:
         randkey = keygen.init_randkey(randkey)
@@ -78,13 +85,18 @@ def adam_unbounded(lossfunc, guess, nsteps=100, learning_rate=0.01,
     opt = optax.adam(learning_rate)
     solver = jaxopt.OptaxSolver(opt=opt, fun=lossfunc, maxiter=nsteps)
     state = solver.init_state(guess, **kwargs)
-    params = [guess]
-    for _ in tqdm.trange(nsteps, desc="Adam Gradient Descent Progress"):
+    params = []
+    params_i = guess
+    for i in tqdm.trange(nsteps, disable=not progress,
+                         desc="Adam Gradient Descent Progress"):
         if randkey is not None:
             randkey, key_i = jax.random.split(randkey)
             kwargs["randkey"] = key_i
-        params_i, state = solver.update(params[-1], state, **kwargs)
-        params.append(params_i)
+        params_i, state = solver.update(params_i, state, **kwargs)
+        if i == nsteps - 1 or (thin and i % thin == thin - 1):
+            params.append(params_i)
+    if not thin:
+        params = params[-1]
 
     return jnp.array(params)
 

diffopt/multigrad/adam.py

@@ -25,12 +25,12 @@
     N_RANKS = 1
 
 
-def trange_no_tqdm(n, desc=None):
+def trange_no_tqdm(n, desc=None, disable=False):
     return range(n)
 
 
-def trange_with_tqdm(n, desc="Adam Gradient Descent Progress"):
-    return tqdm.trange(n, desc=desc)
+def trange_with_tqdm(n, desc="Adam Gradient Descent Progress", disable=False):
+    return tqdm.trange(n, desc=desc, disable=disable)
 
 
 adam_trange = trange_no_tqdm if tqdm is None else trange_with_tqdm
@@ -49,27 +49,32 @@ def _master_wrapper(params, logloss_and_grad_fn, data, randkey=None):
     return loss, grad
 
 
-def _adam_optimizer(params, fn, fn_data, nsteps, learning_rate, randkey=None):
+def _adam_optimizer(params, fn, fn_data, nsteps, learning_rate, randkey=None,
+                    thin=1, progress=True):
     kwargs = {}
     # Note: Might be recommended to use optax instead of jax.example_libraries
     opt_init, opt_update, get_params = jax_opt.adam(learning_rate)
     opt_state = opt_init(params)
 
-    param_steps = [params]
-    for step in adam_trange(nsteps):
+    param_steps = []
+    for step in adam_trange(nsteps, disable=not progress):
         if randkey is not None:
             randkey, key_i = jax.random.split(randkey)
             kwargs["randkey"] = key_i
         _, grad = fn(params, *fn_data, **kwargs)
         opt_state = opt_update(step, grad, opt_state)
         params = get_params(opt_state)
-        param_steps.append(params)
+        if step == nsteps - 1 or (thin and step % thin == thin - 1):
+            param_steps.append(params)
+    if not thin:
+        param_steps = param_steps[-1]
 
     return jnp.array(param_steps)
 
 
 def run_adam_unbounded(logloss_and_grad_fn, params, data, nsteps=100,
-                       learning_rate=0.01, randkey=None):
+                       learning_rate=0.01, randkey=None,
+                       thin=1, progress=True):
     """Run the adam optimizer on a loss function with a custom gradient.
 
     Parameters
@@ -88,6 +93,11 @@ def run_adam_unbounded(logloss_and_grad_fn, params, data, nsteps=100,
     randkey : int | PRNG Key
         If given, a new PRNG Key will be generated at each iteration and be
         passed to `logloss_and_grad_fn` under the "randkey" kwarg
+    thin : int, optional
+        Return parameters for every `thin` iterations, by default 1. Set
+        `thin=0` to only return final parameters
+    progress : bool, optional
+        Display tqdm progress bar, by default True
 
     Returns
     -------
@@ -104,7 +114,7 @@ def run_adam_unbounded(logloss_and_grad_fn, params, data, nsteps=100,
         fn_data = (logloss_and_grad_fn, data)
 
         params = _adam_optimizer(params, fn, fn_data, nsteps, learning_rate,
-                                 randkey=randkey)
+                                 randkey=randkey, thin=thin, progress=progress)
 
         if COMM is not None:
             COMM.bcast("exit", root=0)
@@ -131,7 +141,7 @@ def run_adam_unbounded(logloss_and_grad_fn, params, data, nsteps=100,
 
 
 def run_adam(logloss_and_grad_fn, params, data, nsteps=100, param_bounds=None,
-             learning_rate=0.01, randkey=None):
+             learning_rate=0.01, randkey=None, thin=1, progress=True):
     """Run the adam optimizer on a loss function with a custom gradient.
 
     Parameters
@@ -153,6 +163,11 @@ def run_adam(logloss_and_grad_fn, params, data, nsteps=100, param_bounds=None,
     randkey : int | PRNG Key
         If given, a new PRNG Key will be generated at each iteration and be
         passed to `logloss_and_grad_fn` under the "randkey" kwarg
+    thin : int, optional
+        Return parameters for every `thin` iterations, by default 1. Set
+        `thin=0` to only return final parameters
+    progress : bool, optional
+        Display tqdm progress bar, by default True
 
     Returns
     -------
@@ -162,7 +177,8 @@ def run_adam(logloss_and_grad_fn, params, data, nsteps=100, param_bounds=None,
     if param_bounds is None:
         return run_adam_unbounded(
             logloss_and_grad_fn, params, data, nsteps=nsteps,
-            learning_rate=learning_rate, randkey=randkey)
+            learning_rate=learning_rate, randkey=randkey,
+            thin=thin, progress=progress)
 
     assert len(params) == len(param_bounds)
     if hasattr(param_bounds, "tolist"):
@@ -182,7 +198,8 @@ def unbound_loss_and_grad(uparams, *args, **kwargs):
 
     uparams = apply_trans(params)
     final_uparams = run_adam_unbounded(
-        unbound_loss_and_grad, uparams, data, nsteps, learning_rate, randkey)
+        unbound_loss_and_grad, uparams, data, nsteps, learning_rate, randkey,
+        thin, progress)
 
     if RANK == 0:
         final_params = invert_trans(final_uparams.T).T

diffopt/multigrad/bfgs.py

@@ -18,19 +18,19 @@
     N_RANKS = 1
 
 
-def trange_no_tqdm(n, desc=None):
+def trange_no_tqdm(n, desc=None, disable=False):
     return range(n)
 
 
-def trange_with_tqdm(n, desc="BFGS Gradient Descent Progress"):
-    return tqdm.trange(n, desc=desc, leave=True)
+def trange_with_tqdm(n, desc="BFGS Gradient Descent Progress", disable=False):
+    return tqdm.trange(n, desc=desc, leave=True, disable=disable)
 
 
 bfgs_trange = trange_no_tqdm if tqdm is None else trange_with_tqdm
 
 
 def run_bfgs(loss_and_grad_fn, params, maxsteps=100, param_bounds=None,
-             randkey=None, comm=COMM):
+             randkey=None, progress=True, comm=COMM):
     """Run the adam optimizer on a loss function with a custom gradient.
 
     Parameters
@@ -46,6 +46,8 @@ def run_bfgs(loss_and_grad_fn, params, maxsteps=100, param_bounds=None,
         `None` as the bound for each unbounded parameter, by default None
     randkey : int | PRNG Key (default=None)
         This will be passed to `logloss_and_grad_fn` under the "randkey" kwarg
+    progress : bool, optional
+        Display tqdm progress bar, by default True
     comm : MPI Communicator (default=COMM_WORLD)
         Communicator between all desired MPI ranks
 
@@ -66,7 +68,7 @@ def run_bfgs(loss_and_grad_fn, params, maxsteps=100, param_bounds=None,
         kwargs["randkey"] = randkey
 
     if comm is None or comm.rank == 0:
-        pbar = bfgs_trange(maxsteps)
+        pbar = bfgs_trange(maxsteps, disable=not progress)
 
         # Wrap loss_and_grad function with commands to the worker ranks
         def loss_and_grad_fn_root(params):