Draft: Refactor links and distributions

src/rolch/abc.py

@@ -1,12 +1,14 @@
 from abc import ABC, abstractmethod
-from typing import Tuple, Union
+from typing import Dict, List, Tuple, Union
 
 import numpy as np
 
 
 class LinkFunction(ABC):
     """The base class for the link functions."""
 
+    _valid_structures: List[str]
+
     @abstractmethod
     def link(self, x: np.ndarray) -> np.ndarray:
         """Calculate the Link"""
@@ -22,6 +24,18 @@ def derivative(self, x: np.ndarray) -> np.ndarray:
 
 class Distribution(ABC):
 
+    links: Dict[int, LinkFunction] | List[LinkFunction]
+    _param_structure: Dict[int, str]
+
+    def _check_links(self):
+        for p in range(self.n_params):
+            if self.param_structure[p] not in self.links[p]._valid_structures:
+                raise ValueError(
+                    f"Link function does not match parameter structure for parameter {p}. \n"
+                    f"Parameter structure is {self.param_structure[p]}. \n"
+                    f"Link function supports {self.links[p]._valid_structures}"
+                )
+
     @abstractmethod
     def theta_to_params(self, theta: np.ndarray) -> Tuple:
         """Take the fitted values and return tuple of vectors for distribution parameters."""

src/rolch/distributions/johnsonsu.py

@@ -34,6 +34,8 @@ def __init__(
             self.shape_link,  # skew
             self.tail_link,  # tail
         ]
+        self._param_structure = {0: "vector", 1: "vector", 2: "vector", 3: "vector"}
+        self._check_links()
 
     def theta_to_params(self, theta):
         mu = theta[:, 0]

src/rolch/distributions/normal.py

@@ -13,6 +13,8 @@ def __init__(self, loc_link=IdentityLink(), scale_link=LogLink()):
         self.loc_link = loc_link
         self.scale_link = scale_link
         self.links = [self.loc_link, self.scale_link]
+        self._param_structure = {0: "vector", 1: "vector"}
+        self._check_links()
 
     def theta_to_params(self, theta):
         mu = theta[:, 0]

src/rolch/distributions/studentt.py

@@ -17,6 +17,8 @@ def __init__(
         self.scale_link = scale_link
         self.tail_link = tail_link
         self.links = [self.loc_link, self.scale_link, self.tail_link]
+        self._param_structure = {0: "vector", 1: "vector", 2: "vector"}
+        self._check_links()
 
     def theta_to_params(self, theta):
         mu = theta[:, 0]

src/rolch/link.py

@@ -15,7 +15,7 @@ class LogLink(LinkFunction):
     """
 
     def __init__(self):
-        pass
+        self._valid_structures = ["vector", "matrix", "square_matrix"]
 
     def link(self, x):
         return np.log(np.fmax(x, LOG_LOWER_BOUND))
@@ -38,7 +38,7 @@ class IdentityLink(LinkFunction):
     """
 
     def __init__(self):
-        pass
+        self._valid_structures = ["vector", "matrix", "square_matrix"]
 
     def link(self, x):
         return x
@@ -62,6 +62,7 @@ class LogShiftValueLink(LinkFunction):
 
     def __init__(self, value):
         self.value = value
+        self._valid_structures = ["vector", "matrix", "square_matrix"]
 
     def link(self, x):
         return np.log(x - self.value + LOG_LOWER_BOUND)
@@ -97,7 +98,7 @@ class SqrtLink(LinkFunction):
     """
 
     def __init__(self):
-        pass
+        self._valid_structures = ["vector", "matrix", "square_matrix"]
 
     def link(self, x):
         return np.sqrt(x)
@@ -121,6 +122,7 @@ class SqrtShiftValueLink(LinkFunction):
 
     def __init__(self, value):
         self.value = value
+        self._valid_structures = ["vector", "matrix", "square_matrix"]
 
     def link(self, x):
         return np.sqrt(x - self.value + LOG_LOWER_BOUND)
@@ -156,7 +158,7 @@ class LogIdentLink(LinkFunction):
     """
 
     def __init__(self):
-        pass
+        self._valid_structures = ["vector", "matrix", "square_matrix"]
 
     def link(self, x: np.ndarray):
         return np.where(x <= 1, np.log(x), x - 1)