Check membership with dicts; better default handling for None sentinels

Author: pchlenski

manify/curvature_estimation/_pipelines.py

@@ -35,10 +35,8 @@ def distortion_pipeline(
         A function f(signature) → loss, where signature is a list
         of (curvature, dim) tuples.
     """
-    if embedder_init_kwargs is None:
-        embedder_init_kwargs = {}
-    if embedder_fit_kwargs is None:
-        embedder_fit_kwargs = {}
+    embedder_init_kwargs = embedder_init_kwargs or {}
+    embedder_fit_kwargs = embedder_fit_kwargs or {}
 
     dists = dists.to(pm.device)
     dists_rescaled = dists / dists.max()
@@ -82,14 +80,10 @@ def classifier_pipeline(
     Returns:
         The loss of the classifier on the test set after embedding the distances using the product manifold.
     """
-    if embedder_init_kwargs is None:
-        embedder_init_kwargs = {}
-    if embedder_fit_kwargs is None:
-        embedder_fit_kwargs = {}
-    if model_init_kwargs is None:
-        model_init_kwargs = {}
-    if model_fit_kwargs is None:
-        model_fit_kwargs = {}
+    embedder_init_kwargs = embedder_init_kwargs or {}
+    embedder_fit_kwargs = embedder_fit_kwargs or {}
+    model_init_kwargs = model_init_kwargs or {}
+    model_fit_kwargs = model_fit_kwargs or {}
 
     dists = dists.to(pm.device)
     dists_rescaled = dists / dists.max()

manify/embedders/_base.py

@@ -33,7 +33,7 @@ class BaseEmbedder(BaseEstimator, TransformerMixin, ABC):
     def __init__(self, pm: ProductManifold, random_state: int | None = None, device: str | None = None) -> None:
         self.pm = pm
         self.random_state = random_state
-        self.device = pm.device if device is None else device
+        self.device = device or pm.device
         self.loss_history_: dict[str, list[float]] = {}
         self.is_fitted_: bool = False
 

manify/manifolds.py

@@ -224,14 +224,11 @@ def sample(
             x: Tensor of sampled points on the manifold
             v: Tensor of tangent vectors
         """
-        if z_mean is None:
-            z_mean = self.mu0
+        z_mean = self.mu0 if z_mean is None else z_mean
         z_mean = torch.Tensor(z_mean).reshape(-1, self.ambient_dim).to(self.device)
         n = z_mean.shape[0]
-        if sigma is None:
-            sigma = torch.stack([torch.eye(self.dim)] * n).to(self.device)
-        else:
-            sigma = torch.Tensor(sigma).reshape(-1, self.dim, self.dim).to(self.device)
+        sigma = torch.stack([torch.eye(self.dim)] * n).to(self.device) if sigma is None else sigma
+        sigma = torch.Tensor(sigma).reshape(-1, self.dim, self.dim).to(self.device)
         assert sigma.shape == (
             n,
             self.dim,
@@ -284,14 +281,11 @@ def log_likelihood(
                 `mu` and covariance `sigma`.
         """
         # Default to mu=self.mu0 and sigma=I
-        if mu is None:
-            mu = self.mu0
+        mu = self.mu0 if mu is None else mu
         mu = torch.Tensor(mu).reshape(-1, self.ambient_dim).to(self.device)
         n = mu.shape[0]
-        if sigma is None:
-            sigma = torch.stack([torch.eye(self.dim)] * n).to(self.device)
-        else:
-            sigma = torch.Tensor(sigma).reshape(-1, self.dim, self.dim).to(self.device)
+        sigma = torch.stack([torch.eye(self.dim)] * n).to(self.device) if sigma is None else sigma
+        sigma = torch.Tensor(sigma).reshape(-1, self.dim, self.dim).to(self.device)
 
         # Euclidean case is regular old Gaussian log-likelihood
         if self.type == "E":
@@ -336,8 +330,7 @@ def logmap(
         Returns:
             logmap_result: Tensor representing the result of the logarithmic map from `base` to `x` on the manifold.
         """
-        if base is None:
-            base = self.mu0
+        base = self.mu0 if base is None else base
         return self.manifold.logmap(x=base, y=x)
 
     def expmap(
@@ -355,8 +348,7 @@ def expmap(
         Returns:
             expmap_result: Tensor representing the result of the exponential map applied to `u` at the base point.
         """
-        if base is None:
-            base = self.mu0
+        base = self.mu0 if base is None else base
         return self.manifold.expmap(x=base, u=u)
 
     def stereographic(self, *points: Float[torch.Tensor, "n_points n_dim"]) -> tuple[Manifold, ...]:
@@ -633,18 +625,17 @@ def sample(
             x: Tensor of sampled points on the manifold
             v: Tensor of tangent vectors
         """
-        if z_mean is None:
-            z_mean = self.mu0
+        z_mean = self.mu0 if z_mean is None else z_mean
         z_mean = torch.Tensor(z_mean).reshape(-1, self.ambient_dim).to(self.device)
         n = z_mean.shape[0]
 
-        if sigma_factorized is None:
-            sigma_factorized = [torch.stack([torch.eye(M.dim)] * n) for M in self.P]
-        else:
-            sigma_factorized = [
-                torch.Tensor(sigma).reshape(-1, M.dim, M.dim).to(self.device)
-                for M, sigma in zip(self.P, sigma_factorized, strict=False)
-            ]
+        sigma_factorized = (
+            [torch.stack([torch.eye(M.dim)] * n) for M in self.P] if sigma_factorized is None else sigma_factorized
+        )
+        sigma_factorized = [
+            torch.Tensor(sigma).reshape(-1, M.dim, M.dim).to(self.device)
+            for M, sigma in zip(self.P, sigma_factorized, strict=False)
+        ]
 
         assert all(sigma.shape == (n, M.dim, M.dim) for M, sigma in zip(self.P, sigma_factorized, strict=False)), (
             "Sigma matrices must match the dimensions of the manifolds."
@@ -684,12 +675,12 @@ def log_likelihood(
                 `mu` and covariance `sigma`.
         """
         n = z.shape[0]
-        if mu is None:
-            mu = torch.vstack([self.mu0] * n).to(self.device)
+        mu = torch.vstack([self.mu0] * n).to(self.device) if mu is None else mu
 
-        if sigma_factorized is None:
-            sigma_factorized = [torch.stack([torch.eye(M.dim)] * n) for M in self.P]
-            # Note that this factorization assumes block-diagonal covariance matrices
+        sigma_factorized = (
+            [torch.stack([torch.eye(M.dim)] * n) for M in self.P] if sigma_factorized is None else sigma_factorized
+        )
+        # Note that this factorization assumes block-diagonal covariance matrices
 
         mu_factorized = self.factorize(mu)
         z_factorized = self.factorize(z)
@@ -807,10 +798,8 @@ def gaussian_mixture(
             torch.manual_seed(seed)
 
         # Deal with clusters
-        if num_clusters is None:
-            num_clusters = num_classes
-        else:
-            assert num_clusters >= num_classes, "Number of clusters must be at least as large as number of classes."
+        num_clusters = num_clusters or num_classes
+        assert num_clusters >= num_classes, "Number of clusters must be at least as large as number of classes."
 
         # Adjust covariance matrices for number of dimensions
         if adjust_for_dims:

manify/predictors/_base.py

@@ -47,7 +47,7 @@ def __init__(
         self.pm = pm
         self.task = task
         self.random_state = random_state
-        self.device = pm.device if device is None else device
+        self.device = device or pm.device
         self.loss_history_: dict[str, list[float]] = {}
         self.is_fitted_: bool = False
 

manify/predictors/_kernel.py

@@ -28,8 +28,7 @@ def compute_kernel_and_norm_manifold(
         kernel_matrix: The kernel matrix between source and target points.
         norm_constant: Scalar normalization constant for the kernel.
     """
-    if X_target is None:
-        X_target = X_source
+    X_target = X_source if X_target is None else X_target
 
     ip = manifold.inner(X_source, X_target)
     ip *= manifold.scale
@@ -77,9 +76,7 @@ def product_kernel(
         kernel_matrices: List of kernel matrices for each component manifold.
         norm_constants: List of normalization constants for each kernel.
     """
-    # If X_target is None, set it to X_source
-    if X_target is None:
-        X_target = X_source
+    X_target = X_source if X_target is None else X_target
 
     # Compute the kernel matrix and norm for each manifold
     Ks = []

manify/predictors/decision_tree.py

@@ -292,10 +292,7 @@ def __init__(
 
         # Store hyperparameters
         self.pm = pm
-        if max_depth is None:
-            self.max_depth = -1  # This runs forever since the loop checks depth == 0
-        else:
-            self.max_depth = max_depth
+        self.max_depth = max_depth or -1
         self.min_samples_leaf = min_samples_leaf
         self.min_samples_split = min_samples_split
         self.min_impurity_decrease = min_impurity_decrease
@@ -370,7 +367,7 @@ def _preprocess(
             dims = self.pm.man2dim[i]
 
             # Non-Euclidean manifolds use angular projections
-            if M.type in ["H", "S"]:
+            if M.type in {"H", "S"}:
                 if self.n_features == "d":
                     dim = dims[0]
                     num = X[:, dim : dim + 1]
@@ -516,7 +513,7 @@ def _aggregate_special_dims(
     ) -> tuple[Float[torch.Tensor, "batch ambient_dim"], ProductManifold]:
         special_dims = []
         for i, M in enumerate(self.pm.P):
-            if M.type in ["H", "S"]:
+            if M.type in {"H", "S"}:
                 dim = self.pm.man2dim[i][0]
                 special_dims.append(X[:, dim : dim + 1])
         if len(special_dims) > 0:
@@ -655,10 +652,7 @@ def __init__(
         tree_kwargs: Dict[str, Any] = {}
         self.pm = tree_kwargs["pm"] = pm
         self.task = tree_kwargs["task"] = task
-        if max_depth is None:
-            self.max_depth = tree_kwargs["max_depth"] = -1
-        else:
-            self.max_depth = tree_kwargs["max_depth"] = max_depth
+        self.max_depth = tree_kwargs["max_depth"] = max_depth or -1
         self.min_samples_leaf = tree_kwargs["min_samples_leaf"] = min_samples_leaf
         self.min_samples_split = tree_kwargs["min_samples_split"] = min_samples_split
         self.min_impurity_decrease = tree_kwargs["min_impurity_decrease"] = min_impurity_decrease

manify/predictors/nn/layers.py

@@ -35,16 +35,11 @@ def __init__(
     ):
         super().__init__()
 
-        # Parameters are Euclidean, straightforardly
-        # self.W = torch.rand(in_features, out_features)
+        # Parameters are Euclidean, straightforwardly
         self.W = torch.nn.Parameter(torch.randn(in_features, out_features) * 0.01)
-        # self.b = torch.nn.Parameter(torch.rand(out_features))
 
-        # Noninearity must be applied via the manifold
-        if nonlinearity is None:
-            self.sigma = lambda x: x
-        else:
-            self.sigma = lambda x: manifold.expmap(nonlinearity(manifold.logmap(x)))
+        # Nonlinearity must be applied via the manifold
+        self.sigma = manifold.apply(nonlinearity) if nonlinearity else lambda x: x
 
         # Also store manifold
         self.manifold = manifold

manify/predictors/perceptron.py

@@ -52,11 +52,8 @@ def __init__(
         self.pm = pm  # ProductManifold instance
         self.max_epochs = max_epochs
         self.patience = patience  # Number of consecutive epochs without improvement to consider convergence
-        if weights is None:
-            self.weights = torch.ones(len(pm.P), dtype=torch.float32)
-        else:
-            assert len(weights) == len(pm.P), "Number of weights must match the number of manifolds."
-            self.weights = weights
+        self.weights = torch.ones(len(pm.P), dtype=torch.float32) if weights is None else weights
+        assert len(self.weights) == len(pm.P), "Number of weights must match the number of manifolds."
 
     def fit(
         self, X: Float[torch.Tensor, "n_samples n_manifolds"], y: Int[torch.Tensor, "n_samples"]

manify/predictors/svm.py

@@ -82,7 +82,7 @@ def __init__(
         self.eps = epsilon
         self.task = task
         self.weights = torch.ones(len(pm.P), dtype=torch.float32) if weights is None else weights
-        assert len(self.weights) == len(pm.P), "Number of weights must match manifolds."
+        assert len(self.weights) == len(pm.P), "Number of weights must match the number of manifolds."
 
     def fit(
         self,
@@ -99,7 +99,6 @@ def fit(
             self: Fitted ProductSpaceSVM instance.
         """
         # unique classes
-        # self.classes_ = torch.unique(y).tolist()
         self._store_classes(y)
         n = X.shape[0]
 

manify/utils/benchmarks.py

@@ -61,15 +61,14 @@ def _score(
     use_torch: bool = False,
     score: list[SCORETYPE] | None = None,
 ) -> dict[SCORETYPE, float]:
-    if score is None:
-        score = ["accuracy", "f1-micro"]
-    if y_pred_override is not None:
-        y_pred = y_pred_override
-    else:
-        assert model is not None, "Model must be provided if y_pred_override is not given"
-        y_pred = model.predict(_X)
+    """Helper function to score a model."""
+    score = score or ["accuracy", "f1-micro"]
+    assert model is not None or y_pred_override is not None, "Model must be provided if y_pred_override is not given"
+    y_pred = y_pred_override if y_pred_override is not None else model.predict(_X)  # type: ignore
+
     if use_torch:
         y_pred = y_pred.detach().cpu().numpy()
+
     scoring_funcs = {
         "accuracy": accuracy_score,
         "f1-micro": lambda y, p: f1_score(y, p, average="micro"),
@@ -164,42 +163,40 @@ def benchmark(
     Returns:
         Dictionary mapping model names to their corresponding evaluation scores.
     """
-    if score is None:
-        score = ["accuracy", "f1-micro", "f1-macro"]
-    if models is None:
-        models = [
-            "sklearn_dt",
-            "sklearn_rf",
-            "product_dt",
-            "product_rf",
-            "tangent_dt",
-            "tangent_rf",
-            "knn",
-            "ps_perceptron",
-            # "svm",
-            "ps_svm",
-            # "tangent_mlp",
-            "ambient_mlp",
-            "tangent_gcn",
-            "ambient_gcn",
-            "kappa_gcn",
-            "ambient_mlr",
-            "tangent_mlr",
-            "kappa_mlr",
-            "single_manifold_rf",
-        ]
+    score = score or ["accuracy", "f1-micro", "f1-macro"]
+    models = models or [
+        "sklearn_dt",
+        "sklearn_rf",
+        "product_dt",
+        "product_rf",
+        "tangent_dt",
+        "tangent_rf",
+        "knn",
+        "ps_perceptron",
+        "svm",
+        "ps_svm",
+        "tangent_mlp",
+        "ambient_mlp",
+        "tangent_gcn",
+        "ambient_gcn",
+        "kappa_gcn",
+        "ambient_mlr",
+        "tangent_mlr",
+        "kappa_mlr",
+        "single_manifold_rf",
+    ]
 
     # Input validation on (task, score) pairing
-    if task in ["classification", "link_prediction"]:
-        assert all(s in ["accuracy", "f1-micro", "f1-macro", "time"] for s in score)
+    if task in {"classification", "link_prediction"}:
+        assert all(s in {"accuracy", "f1-micro", "f1-macro", "time"} for s in score)
     elif task == "regression":
-        assert all(s in ["mse", "rmse", "percent_rmse", "time"] for s in score)
+        assert all(s in {"mse", "rmse", "percent_rmse", "time"} for s in score)
 
     # Input validation on (task, score) pairing
-    if task in ["classification", "link_prediction"]:
-        assert all(s in ["accuracy", "f1-micro", "f1-macro", "time"] for s in score)
+    if task in {"classification", "link_prediction"}:
+        assert all(s in {"accuracy", "f1-micro", "f1-macro", "time"} for s in score)
     elif task == "regression":
-        assert all(s in ["mse", "rmse", "percent_rmse", "time"] for s in score)
+        assert all(s in {"mse", "rmse", "percent_rmse", "time"} for s in score)
     else:
         raise ValueError(f"Unknown task: {task}")
 
@@ -243,7 +240,7 @@ def benchmark(
         X_train, X_test, y_train, y_test, train_idx, test_idx = train_test_split(X, y, np.arange(len(X)), test_size=0.2)
 
     # Make sure classification labels are formatted correctly
-    if task in ["classification", "link_prediction"]:
+    if task in {"classification", "link_prediction"}:
         y = torch.unique(y, return_inverse=True)[1]
         y_train = y[train_idx]
         y_test = y[test_idx]
@@ -302,7 +299,7 @@ def benchmark(
     nn_train_kwargs = {"epochs": epochs, "lr": lr}
 
     # Define your models
-    if task in ["classification", "link_prediction"]:
+    if task in {"classification", "link_prediction"}:
         dt_class = DecisionTreeClassifier
         rf_class = RandomForestClassifier
         knn_class = KNeighborsClassifier