Fix root date issues to unpin numpyro

.devcontainer/Dockerfile

@@ -0,0 +1,21 @@
+# See here for image contents: https://github.com/microsoft/vscode-dev-containers/tree/v0.236.0/containers/python-3/.devcontainer/base.Dockerfile
+
+# [Choice] Python version (use -bullseye variants on local arm64/Apple Silicon): 3, 3.10, 3.9, 3.8, 3.7, 3.6, 3-bullseye, 3.10-bullseye, 3.9-bullseye, 3.8-bullseye, 3.7-bullseye, 3.6-bullseye, 3-buster, 3.10-buster, 3.9-buster, 3.8-buster, 3.7-buster, 3.6-buster
+ARG VARIANT="3.10-bullseye"
+FROM mcr.microsoft.com/vscode/devcontainers/python:0-${VARIANT}
+
+# [Choice] Node.js version: none, lts/*, 16, 14, 12, 10
+ARG NODE_VERSION="none"
+RUN if [ "${NODE_VERSION}" != "none" ]; then su vscode -c "umask 0002 && . /usr/local/share/nvm/nvm.sh && nvm install ${NODE_VERSION} 2>&1"; fi
+
+# [Optional] If your pip requirements rarely change, uncomment this section to add them to the image.
+# COPY requirements.txt /tmp/pip-tmp/
+# RUN pip3 --disable-pip-version-check --no-cache-dir install -r /tmp/pip-tmp/requirements.txt \
+#    && rm -rf /tmp/pip-tmp
+
+# [Optional] Uncomment this section to install additional OS packages.
+# RUN apt-get update && export DEBIAN_FRONTEND=noninteractive \
+#     && apt-get -y install --no-install-recommends <your-package-list-here>
+
+# [Optional] Uncomment this line to install global node packages.
+# RUN su vscode -c "source /usr/local/share/nvm/nvm.sh && npm install -g <your-package-here>" 2>&1

.devcontainer/devcontainer.json

@@ -0,0 +1,57 @@
+// For format details, see https://aka.ms/devcontainer.json. For config options, see the README at:
+// https://github.com/microsoft/vscode-dev-containers/tree/v0.236.0/containers/python-3
+{
+	"name": "Python 3",
+	"build": {
+		"dockerfile": "Dockerfile",
+		"context": "..",
+		"args": { 
+			// Update 'VARIANT' to pick a Python version: 3, 3.10, 3.9, 3.8, 3.7, 3.6
+			// Append -bullseye or -buster to pin to an OS version.
+			// Use -bullseye variants on local on arm64/Apple Silicon.
+			"VARIANT": "3.10-bullseye",
+			// Options
+			"NODE_VERSION": "lts/*"
+		}
+	},
+
+	// Configure tool-specific properties.
+	"customizations": {
+		// Configure properties specific to VS Code.
+		"vscode": {
+			// Set *default* container specific settings.json values on container create.
+			"settings": { 
+				"python.defaultInterpreterPath": "/usr/local/bin/python",
+				"python.linting.enabled": true,
+				"python.linting.pylintEnabled": true,
+				"python.formatting.autopep8Path": "/usr/local/py-utils/bin/autopep8",
+				"python.formatting.blackPath": "/usr/local/py-utils/bin/black",
+				"python.formatting.yapfPath": "/usr/local/py-utils/bin/yapf",
+				"python.linting.banditPath": "/usr/local/py-utils/bin/bandit",
+				"python.linting.flake8Path": "/usr/local/py-utils/bin/flake8",
+				"python.linting.mypyPath": "/usr/local/py-utils/bin/mypy",
+				"python.linting.pycodestylePath": "/usr/local/py-utils/bin/pycodestyle",
+				"python.linting.pydocstylePath": "/usr/local/py-utils/bin/pydocstyle",
+				"python.linting.pylintPath": "/usr/local/py-utils/bin/pylint"
+			},
+
+			// Add the IDs of extensions you want installed when the container is created.
+			"extensions": [
+				"ms-python.python",
+				"ms-python.vscode-pylance"
+			]
+		}
+	},
+
+	// Use 'forwardPorts' to make a list of ports inside the container available locally.
+	// "forwardPorts": [],
+
+	// Use 'postCreateCommand' to run commands after the container is created.
+	// "postCreateCommand": "pip3 install --user -r requirements.txt",
+
+	// Comment out to connect as root instead. More info: https://aka.ms/vscode-remote/containers/non-root.
+	"remoteUser": "vscode",
+	"features": {
+		"git": "os-provided"
+	}
+}

setup.cfg

@@ -20,7 +20,7 @@ package_dir =
 packages = find:
 python_requires = >=3.6
 install_requires =
-    numpyro==0.9.0
+    numpyro
     alive-progress
     pandas
     treeswift

src/chronumental/__main__.py

@@ -378,7 +378,7 @@ def main():
         print(f"Wrote tree to {args.tree_out}")
 
         origin_date = lookup[reference_point][0]
-        output_dates = {name: origin_date +  datetime.timedelta(days=(x + params['root_date'].tolist())) for name,x in total_lengths_in_time.items()}
+        output_dates = {name: origin_date +  datetime.timedelta(days=(x + params['root_date_mu'].tolist())) for name,x in total_lengths_in_time.items()}
 
         names, values = zip(*output_dates.items())
         output_meta = pd.DataFrame({"strain": names, "predicted_date": values})

src/chronumental/models.py

@@ -20,7 +20,7 @@ def __init__(self,**kwargs):
     def get_logging_results(self,params):
         results = collections.OrderedDict()
         times = self.get_branch_times(params)
-        new_dates = self.calc_dates(times, params['root_date'])
+        new_dates = self.calc_dates(times, params['root_date_mu'])
         results['date_cor'] = onp.corrcoef(
             self.terminal_target_dates_array,
             new_dates)[0, 1]
@@ -37,7 +37,8 @@ def get_logging_results(self,params):
         results['length_cor'] = onp.corrcoef(
             self.branch_distances_array,
             times)[0, 1]  # This correlation should be relatively high
-        results['root_date'] = params['root_date']
+
+        results['root_date'] = params['root_date_mu']
         return results
 
 
@@ -107,7 +108,9 @@ def model(self):
 
 
     def guide(self):
-        root_date = numpyro.param("root_date", -365*self.ref_point_distance/self.clock_rate) 
+        root_date_mu = numpyro.param("root_date_mu", -365*self.ref_point_distance/self.clock_rate) 
+
+        root_date = numpyro.sample("root_date",dist.Delta(root_date_mu))
 
         time_length_mu = numpyro.param("time_length_mu", self.initial_time,
                                 constraint=dist.constraints.positive)