no warning for normalization if d_method was "fractal"

Author: katosh

CHANGELOG.md

@@ -1,3 +1,11 @@
+# v1.6.1
+
+ - improve normalization warning to only show when appropriate
+ - introduce "manual" d_method for explicitly provided d values
+ - include detailed information in normalization warnings and messages
+ - store d and d_method in predictors for better warning behavior
+ - automatically set d_method to "manual" when d is explicitly provided
+
 # v1.6.0
 
  - use [PyNNDescent](https://github.com/lmcinnes/pynndescent) for faster nearest neighbor distance computation

mellon/base_predictor.py

@@ -92,13 +92,27 @@ class Predictor(ABC):
     n_input_features : int
         The number of features/dimensions of the cell-state representation the predictor was
         trained on. This is used for validation of input data.
+        
+    d : int or None
+        The intrinsic dimensionality of the data used to create this predictor.
+        Only stored to provide appropriate warnings for normalization.
+        
+    d_method : str or None
+        The method used to compute the intrinsic dimensionality of the data.
+        Only stored to provide appropriate warnings for normalization.
     """
 
     # number of features of input data (x.shape[1]) to be specified in __init__
     n_input_features: int
 
     # number of observations trained on (x.shape[0]) to be specified in __init__
     n_obs: int
+    
+    # intrinsic dimensionality of the data, stored only for warning purposes
+    d: int = None
+    
+    # method used to compute the intrinsic dimensionality, stored only for warning purposes
+    d_method: str = None
 
     # a set of attribute names that should be saved to reconstruct the object
     _state_variables: Union[Set, List]
@@ -210,9 +224,30 @@ def mean(self, x, normalize=False):
                 )
                 logger.error(message)
                 raise ValueError(message)
-            logger.warning(
-                'The normalization is only effective if the density was trained with d_method="fractal".'
-            )
+            # Check conditions for warning about normalization
+            if self.d_method == "fractal":
+                # No warning needed for fractal method
+                pass
+            elif self.d_method == "manual":
+                # For manual d, show info message
+                logger.info(
+                    f"Using normalization with manually set d={self.d}. "
+                    "Note: Normalization is most effective when d approximates the intrinsic dimensionality of the data."
+                )
+            elif self.d_method is None and isinstance(self.d, (int, float)) and float(self.d).is_integer():
+                # For None d_method with integer d, show warning
+                logger.warning(
+                    f"The normalization is only effective if d approximates the intrinsic dimensionality. "
+                    f"Current values: d_method={self.d_method}, d={self.d}. "
+                    f'Consider using d_method="fractal" for more accurate results.'
+                )
+            elif self.d_method == "embedding":
+                # For embedding method, show warning
+                logger.warning(
+                    f"The normalization is only effective if d approximates the intrinsic dimensionality. "
+                    f"Current values: d_method={self.d_method}, d={self.d}. "
+                    f'Consider using d_method="fractal" for more accurate results.'
+                )
             return self._mean(x) - log(self.n_obs)
         else:
             return self._mean(x)
@@ -396,6 +431,8 @@ def __getstate__(self):
             {
                 "n_input_features": self.n_input_features,
                 "n_obs": self.n_obs,
+                "d": self.d,
+                "d_method": self.d_method,
                 "_state_variables": self._state_variables,
             }
         )
@@ -755,9 +792,30 @@ def mean(self, Xnew, time=None, normalize=False):
                 )
                 logger.error(message)
                 raise ValueError(message)
-            logger.warning(
-                'The normalization is only effective if the density was trained with d_method="fractal".'
-            )
+            # Check conditions for warning about normalization
+            if self.d_method == "fractal":
+                # No warning needed for fractal method
+                pass
+            elif self.d_method == "manual":
+                # For manual d, show info message
+                logger.info(
+                    f"Using normalization with manually set d={self.d}. "
+                    "Note: Normalization is most effective when d approximates the intrinsic dimensionality of the data."
+                )
+            elif self.d_method is None and isinstance(self.d, (int, float)) and float(self.d).is_integer():
+                # For None d_method with integer d, show warning
+                logger.warning(
+                    f"The normalization is only effective if d approximates the intrinsic dimensionality. "
+                    f"Current values: d_method={self.d_method}, d={self.d}. "
+                    f'Consider using d_method="fractal" for more accurate results.'
+                )
+            elif self.d_method == "embedding":
+                # For embedding method, show warning
+                logger.warning(
+                    f"The normalization is only effective if d approximates the intrinsic dimensionality. "
+                    f"Current values: d_method={self.d_method}, d={self.d}. "
+                    f'Consider using d_method="fractal" for more accurate results.'
+                )
             return self._mean(Xnew) - log(self.n_obs)
         else:
             return self._mean(Xnew)

mellon/density_estimator.py

@@ -224,9 +224,14 @@ def __init__(
             jit=jit,
             check_rank=check_rank,
         )
-        self.d_method = validate_string(
-            d_method, "d_method", choices={"fractal", "embedding"}
-        )
+        # If d is explicitly provided, set d_method to "manual"
+        if d is not None:
+            self.d_method = "manual"
+            logger.info(f"Explicitly provided d={d}, setting d_method to 'manual'.")
+        else:
+            self.d_method = validate_string(
+                d_method, "d_method", choices={"fractal", "embedding", "manual"}
+            )
         self.transform = None
         self.loss_func = None
         self.opt_state = None
@@ -305,7 +310,12 @@ def _compute_d(self):
         if self.d_method == "fractal":
             d = compute_d_factal(x)
             logger.info(f"Using d={d}.")
+        elif self.d_method == "manual":
+            # For manual method, d is already set, so we don't need to compute it
+            d = self.d
+            logger.info(f"Using manually set d={d}.")
         else:
+            # embedding method uses the number of dimensions
             d = compute_d(x)
             logger.info(
                 f"Using embedding dimensionality d={d}. "
@@ -382,6 +392,10 @@ def _set_log_density_func(self):
             y_is_mean=True,
             with_uncertainty=with_uncertainty,
         )
+        log_density_func.n_obs = self.x.shape[0]
+        # Store d and d_method for warning purposes
+        log_density_func.d = self.d
+        log_density_func.d_method = self.d_method
         self.log_density_func = log_density_func
 
     def prepare_inference(self, x):

mellon/time_sensitive_density_estimator.py

@@ -288,9 +288,14 @@ def __init__(
         if not isinstance(density_estimator_kwargs, dict):
             raise ValueError("density_estimator_kwargs needs to be a dictionary.")
         self.density_estimator_kwargs = density_estimator_kwargs
-        self.d_method = validate_string(
-            d_method, "d_method", choices={"fractal", "embedding"}
-        )
+        # If d is explicitly provided, set d_method to "manual"
+        if d is not None:
+            self.d_method = "manual"
+            logger.info(f"Explicitly provided d={d}, setting d_method to 'manual'.")
+        else:
+            self.d_method = validate_string(
+                d_method, "d_method", choices={"fractal", "embedding", "manual"}
+            )
         self.ls_time = validate_positive_float(ls_time, "ls_time", optional=True)
         self.ls_time_factor = validate_positive_float(ls_time_factor, "ls_time_factor")
         self._save_intermediate_ls_times = _save_intermediate_ls_times
@@ -418,8 +423,17 @@ def _compute_d(self):
             logger.warning("Using EXPERIMENTAL fractal dimensionality selection.")
             d = compute_d_factal(x)
             logger.info(f"Using d={d}.")
+        elif self.d_method == "manual":
+            # For manual method, d is already set, so we don't need to compute it
+            d = self.d
+            logger.info(f"Using manually set d={d}.")
         else:
+            # embedding method uses the number of dimensions
             d = compute_d(x)
+            logger.info(
+                f"Using embedding dimensionality d={d}. "
+                'Use d_method="fractal" to enable effective density normalization.'
+            )
         if d > 50:
             message = f"""The detected dimensionality of the data is over 50,
             which is likely to cause numerical instability issues.
@@ -580,6 +594,9 @@ def _set_log_density_func(self):
             with_uncertainty=with_uncertainty,
         )
         log_density_func.n_obs = compute_average_cell_count(x, normalize)
+        # Store d and d_method for warning purposes
+        log_density_func.d = self.d
+        log_density_func.d_method = self.d_method
         self.log_density_func = log_density_func
 
     def prepare_inference(self, x, times=None):

mellon/version.py

@@ -1,3 +1,3 @@
 """Version information."""
 
-__version__ = "1.6.0"
+__version__ = "1.6.1"