debug score test and add docstring

SpatialDE/_internal/score_test.py

@@ -303,14 +303,14 @@ def _fit_null(self, y: tf.Tensor) -> NullModel:
         res = minimize(
             lambda *args: self._negative_normal_loglik(*args).numpy(),
             x0=[
-                tf.math.log(tf.reduce_mean(scaledy)),
+                tf.reduce_mean(scaledy),
                 tf.zeros_like(y),
             ],
             args=(tf.cast(y, tf.float64), self.sizefactors),
             jac=lambda *args: self._grad_negative_normal_loglik(*args).numpy(),
             method="bfgs",
         )
-        mu = tf.exp(res.x[0]) * self.sizefactors
+        mu = res.x[0] * self.sizefactors
         sigma = tf.exp(res.x[1])
         return self.NullModel(mu, sigma)
 
@@ -331,7 +331,7 @@ def _do_test(
     ) -> Tuple[tf.Tensor, tf.Tensor, tf.Tensor, tf.Tensor]:
         W = 1 / sigma**2 # Delta W^-1
         stat = 0.5 * tf.reduce_sum(
-            (rawy / mu - 1) * W * tf.tensordot(K, W * (rawy / mu - 1), axes=(-1, -1)), axis=-1
+            (rawy - mu) * W * tf.tensordot(K, W * (rawy - mu), axes=(-1, -1)), axis=-1
         )
 
         P = tf.linalg.diag(W) - W[:, tf.newaxis] * W[tf.newaxis, :] / tf.reduce_sum(W)

SpatialDE/de_test.py

@@ -64,8 +64,8 @@ def test(
     kernel_space: Optional[Dict[str, Union[float, List[float]]]] = None,
     sizefactors: Optional[np.ndarray] = None,
     stack_kernels: Optional[bool] = None,
+    obs_dist: Literal["NegativeBinomial", "Normal"] = "NegativeBinomial",
     use_cache: bool = True,
-    obs_dist: Literal["NegativeBinomial", "Normal"] = "NegativeBinomial"
 ) -> Tuple[pd.DataFrame, Union[pd.DataFrame, None]]:
     """
     Test for spatially variable genes.
@@ -96,6 +96,7 @@ def test(
             the kernel matrices. This leads to increased memory consumption, but will drastically improve runtime
             on GPUs for smaller data sets. Defaults to ``True`` for datasets with less than 2000 observations and
             ``False`` otherwise.
+        obs_dist: Distribution of the observations. If set as "Normal", model the regression to have Gaussian mean field error with identity link function.
         use_cache: Whether to use a pre-computed distance matrix for all kernels instead of computing the distance
             matrix anew for each kernel. Increases memory consumption, but is somewhat faster.