TensorFlow 2 model fitting for negative binomial noise

batchglm/api/models/tf2/glm_beta.py

@@ -1,2 +1,2 @@
-#from batchglm.models.glm_beta import InputDataGLM, Model, Simulator
-#from batchglm.train.tf2.glm_beta import Estimator
+from batchglm.models.glm_beta import InputDataGLM, Model, Simulator
+from batchglm.train.tf2.glm_beta import Estimator

batchglm/api/models/tf2/glm_norm.py

@@ -1,2 +1,2 @@
-#from batchglm.models.glm_norm import InputDataGLM, Model, Simulator
-#from batchglm.train.tf2.glm_norm import Estimator
+from batchglm.models.glm_norm import InputDataGLM, Model, Simulator
+from batchglm.train.tf2.glm_norm import Estimator

batchglm/pkg_constants.py

@@ -11,22 +11,34 @@
 
 # Trust region hyper parameters:
 TRUST_REGION_RADIUS_INIT = 100.
+TRUST_REGION_RADIUS_INIT_SCALE = 1.
 TRUST_REGION_ETA0 = 0.
 TRUST_REGION_ETA1 = 0.25
 TRUST_REGION_ETA2 = 0.25
 TRUST_REGION_T1 = 0.5  # Fast collapse to avoid trailing.
 TRUST_REGION_T2 = 1.5  # Allow expansion if not shrinking.
-TRUST_REGION_UPPER_BOUND = 1e5
+TRUST_REGION_UPPER_BOUND = 1e40
 
-TRUST_REGIONT_T1_IRLS_GD_TR_SCALE = 1
+TRUST_REGIONT_T1_IRLS_GD_TR_SCALE = 0.5
+TRUST_REGIONT_T2_IRLS_GD_TR_SCALE = 1.5
 
 # Convergence hyper-parameters:
-LLTOL_BY_FEATURE = 1e-10
+LLTOL_BY_FEATURE = 1e-12
 XTOL_BY_FEATURE_LOC = 1e-8
 XTOL_BY_FEATURE_SCALE = 1e-6
 GTOL_BY_FEATURE_LOC = 1e-8
 GTOL_BY_FEATURE_SCALE = 1e-8
 
+TRTOL_BY_FEATURE_LOC = 1e-8
+TRTOL_BY_FEATURE_SCALE = 1e-6
+
+FEATUREWISE_THRESHOLD = 10  # the minimal number of features to converge before next featurewise batch
+FEATUREWISE_RECALCULATE = False # if set to True, recalculate the results from the previous train step
+
+WOLFE_C1 = 1e-3
+WOLFE_C2 = 0.99
+ALPHA0 = 100
+
 try:
     import tensorflow as tf
 

batchglm/train/tf2/__init__.py

@@ -0,0 +1,3 @@
+from . import glm_nb as nb
+from . import glm_norm as norm
+from . import glm_beta as beta

batchglm/train/tf2/base/__init__.py

@@ -0,0 +1,3 @@
+from .estimator import TFEstimator
+from .model import ProcessModelBase, ModelBase
+from .optim import OptimizerBase

batchglm/train/tf2/base/estimator.py

@@ -0,0 +1,28 @@
+from .model import ModelBase
+import tensorflow as tf
+
+
+class TFEstimator:
+    model: ModelBase
+
+    def __init__(self, input_data, dtype):
+
+        self._input_data = input_data
+        self.dtype = dtype
+
+    def _train(
+            self,
+            is_batched: bool,
+            batch_size: int,
+            optimizer_object: tf.keras.optimizers.Optimizer,
+            convergence_criteria: str,
+            stopping_criteria: int,
+            autograd: bool,
+            featurewise: bool,
+            benchmark: bool,
+            optimizer: str
+    ):
+        pass
+
+    def fetch_fn(self, idx):
+        pass

batchglm/train/tf2/base/external.py

@@ -0,0 +1,4 @@
+#from batchglm.models.base import _Estimator_Base
+#from batchglm.xarray_sparse import SparseXArrayDataArray, SparseXArrayDataSet
+#import batchglm.utils.stats as stat_utils
+from batchglm import pkg_constants

batchglm/train/tf2/base/model.py

@@ -0,0 +1,47 @@
+import abc
+import logging
+import tensorflow as tf
+import numpy as np
+
+logger = logging.getLogger(__name__)
+
+
+class ModelBase(tf.keras.Model, metaclass=abc.ABCMeta):
+
+    def __init__(self, dtype):
+        super(ModelBase, self).__init__(dtype=dtype)
+
+    @abc.abstractmethod
+    def call(self, inputs, training=False, mask=None):
+        pass
+
+
+class ProcessModelBase:
+
+    @abc.abstractmethod
+    def param_bounds(self, dtype):
+        pass
+
+    def tf_clip_param(
+            self,
+            param,
+            name
+    ):
+        bounds_min, bounds_max = self.param_bounds(param.dtype)
+        return tf.clip_by_value(
+            param,
+            bounds_min[name],
+            bounds_max[name]
+        )
+
+    def np_clip_param(
+            self,
+            param,
+            name
+    ):
+        bounds_min, bounds_max = self.param_bounds(param.dtype)
+        return np.clip(
+            param,
+            bounds_min[name],
+            bounds_max[name]
+        )

batchglm/train/tf2/base/optim.py

@@ -0,0 +1,52 @@
+import abc
+import logging
+import tensorflow as tf
+
+logger = logging.getLogger("batchglm")
+
+
+class OptimizerBase(tf.keras.optimizers.Optimizer, metaclass=abc.ABCMeta):
+
+    def __init__(self, name):
+        super(OptimizerBase, self).__init__(name=name)
+
+    @abc.abstractmethod
+    def _resource_apply_dense(self, grad, handle):
+        pass
+
+    @abc.abstractmethod
+    def _resource_apply_sparse(self, grad, handle, apply_state):
+        pass
+
+    @abc.abstractmethod
+    def _create_slots(self):
+        pass
+
+    """
+    @property
+    @abc.abstractmethod
+    def vars(self):
+        pass
+
+    @property
+    @abc.abstractmethod
+    def gradients(self):
+        return None
+
+    @property
+    @abc.abstractmethod
+    def hessians(self):
+        pass
+
+    @property
+    @abc.abstractmethod
+    def fims(self):
+        pass
+
+    @abc.abstractmethod
+    def step(self, learning_rate):
+        pass
+    """
+    @abc.abstractmethod
+    def get_config(self):
+        pass

batchglm/train/tf2/base_glm/README.md

@@ -0,0 +1,2 @@
+# Classes with GLM specific code.
+All noise models that are in the GLM category inherit all of these classes.

batchglm/train/tf2/base_glm/__init__.py

@@ -0,0 +1,12 @@
+from .processModel import ProcessModelGLM
+from .model import GLM
+
+from .estimator import Estimator
+from .vars import ModelVarsGLM
+from .layers import LinearLocGLM, LinearScaleGLM, LinkerLocGLM, LinkerScaleGLM
+from .layers import LikelihoodGLM, UnpackParamsGLM
+from .layers_gradients import JacobianGLM, HessianGLM, FIMGLM
+from .optim import NR, IRLS
+
+from .generator import DataGenerator
+from .convergence import ConvergenceCalculator