Merge pull request #446 from lnccbrown/444-fix-initialization-with-nans

fix key error bug when going through initialization, make safe priors…
lnccbrown · Jun 12, 2024 · 3d2f301 · 3d2f301
2 parents 3f7c518 + 353c446
commit 3d2f301
Show file tree

Hide file tree

Showing 7 changed files with 539 additions and 42 deletions.
diff --git a/src/hssm/datasets.py b/src/hssm/datasets.py
@@ -5,7 +5,7 @@
 """
 
 import os
-from typing import NamedTuple, Optional, Union
+from typing import NamedTuple, Optional
 
 import pandas as pd
 
@@ -34,7 +34,7 @@ class FileMetadata(NamedTuple):
 }
 
 
-def load_data(dataset: Optional[str] = None) -> Union[pd.DataFrame, str]:
+def load_data(dataset: Optional[str] = None) -> pd.DataFrame | str:
     """
     Load a dataset as a pandas DataFrame.
 

diff --git a/src/hssm/defaults.py b/src/hssm/defaults.py
@@ -92,7 +92,6 @@ class DefaultConfig(TypedDict):
                     "t": {
                         "name": "HalfNormal",
                         "sigma": 2.0,
-                        "initval": 0.05,
                     },
                 },
                 "extra_fields": None,
@@ -104,7 +103,6 @@ class DefaultConfig(TypedDict):
                     "t": {
                         "name": "HalfNormal",
                         "sigma": 2.0,
-                        "initval": 0.05,
                     },
                 },
                 "bounds": {
@@ -123,7 +121,6 @@ class DefaultConfig(TypedDict):
                     "t": {
                         "name": "HalfNormal",
                         "sigma": 2.0,
-                        "initval": 0.05,
                     },
                 },
                 "extra_fields": None,
@@ -143,7 +140,6 @@ class DefaultConfig(TypedDict):
                     "t": {
                         "name": "HalfNormal",
                         "sigma": 2.0,
-                        "initval": 0.05,
                     },
                 },
                 "extra_fields": None,
@@ -155,7 +151,6 @@ class DefaultConfig(TypedDict):
                     "t": {
                         "name": "HalfNormal",
                         "sigma": 2.0,
-                        "initval": 0.05,
                     },
                 },
                 "bounds": {
@@ -175,7 +170,6 @@ class DefaultConfig(TypedDict):
                     "t": {
                         "name": "HalfNormal",
                         "sigma": 2.0,
-                        "initval": 0.05,
                     },
                 },
                 "extra_fields": None,
@@ -195,7 +189,6 @@ class DefaultConfig(TypedDict):
                     "t": {
                         "name": "HalfNormal",
                         "sigma": 2.0,
-                        "initval": 0.05,
                     },
                 },
                 "extra_fields": None,
@@ -328,6 +321,7 @@ class DefaultConfig(TypedDict):
     },
 }
 
+# TODO: Initval settings could be specified directly in model config as well.
 INITVAL_SETTINGS = {
     # logit link function case
     # should never use priors with bounds,
@@ -343,7 +337,7 @@ class DefaultConfig(TypedDict):
     },
     # identity link function case,
     # need to take care of_log__ and _interval__ variables
-    "None": {
+    None: {
         "t": 0.025,
         "t_Intercept": 0.025,
         "a": 1.5,