Refactor ganglion cell (#151)

* Add patch reconstruction function to ganglionCell.py

Added a new function to reconstruct patches from input data, improving data handling in the RetinalGanglionCell class.

* Add ratio of Gaussian filter function

Author: Faezehabibi

ngclearn/components/input_encoders/ganglionCell.py

@@ -22,6 +22,14 @@ def _create_dog_filter(patch_shape, sigma, k=1.6, lmbda=1):
     dog = g1 - lmbda * g2
     return dog #- jnp.mean(dog)
 
+
+def _create_ratio_of_gauss_filter(patch_shape, sigma, k=1.6):
+    g1 = _create_gaussian_filter(patch_shape, sigma=sigma)
+    g2 = _create_gaussian_filter(patch_shape, sigma=sigma * k)
+    rog = g1 / (g2 + 1e-8)
+    return rog
+    
+
 def _create_patches(obs, patch_shape, step_shape):
     """
     Extract 2D patches from a batch of images using a sliding window.
@@ -61,6 +69,29 @@ def _create_patches(obs, patch_shape, step_shape):
 
     return patches
 
+def _reconstruct(patches, nx_ny, area_shape, patch_shape, step_shape):
+    # patches: (N, nx * ny, px, py)
+
+    B = len(patches)
+    nx, ny = nx_ny
+    ix, iy = area_shape
+    px, py = patch_shape
+    sx, sy = step_shape
+    x = jnp.zeros((B, ix, iy))
+    counts = jnp.zeros((ix, iy))
+
+    idx = 0
+    for i in range(ny):
+        for j in range(nx):
+            di = i * sx
+            dj = j * sy
+            x = x.at[:, di:di + px, dj:dj + py].add(patches[:, idx])
+            counts = counts.at[di:di + px, dj:dj + py].add(1.0)
+            idx += 1
+
+    return x / counts[None, :, :]
+
+
 
 class RetinalGanglionCell(JaxComponent):
     """
@@ -121,10 +152,18 @@ def __init__(
         self.step_shape = step_shape
 
         _filter = jnp.ones(self.patch_shape)
-        if filter_type == 'gaussian':
+
+        if self.filter_type == 'gaussian':
+            print("filter type is ", self.filter_type)
             _filter = _create_gaussian_filter(patch_shape=self.patch_shape, sigma=self.sigma)
-        elif filter_type == 'difference_of_gaussian':
+            
+        elif self.filter_type in ["difference_of_gaussian", "DoG"]:
+            print("filter type is difference of gaussian: f(x) = p1 - p2")
             _filter = _create_dog_filter(patch_shape=self.patch_shape, sigma=sigma)
+            
+        elif self.filter_type in ["ratio_of_gaussian", "RoG"]:
+            print("filter type is ratio of gaussian: f(x) = p1 / p2")
+            _filter = _create_ratio_of_gauss_filter(patch_shape=self.patch_shape, sigma=sigma)
 
         # ═════════════════ compartments initial values ════════════════════
         in_restVals = jnp.zeros((batch_size, *self.area_shape)) ## input: (B | ix | iy)
@@ -158,27 +197,16 @@ def advance_state(self, t):
 
         self.outputs.set(outputs)
 
+
+
     @compilable
     def reset(self):  ## reset core components/statistics
-        # self.batched_reset(batch_size=self.batch_size)  ## arg = batch_size data-member
         in_restVals = jnp.zeros((self.batch_size, *self.area_shape))      ## input: (B | ix | iy)
         out_restVals = jnp.zeros((self.batch_size,      ## output.shape: (B | n_cells * px * py)
                                   self.n_cells * self.patch_shape[0] * self.patch_shape[1]))
         self.inputs.set(in_restVals)
         self.outputs.set(out_restVals)
 
-    # Viet: NOTE: we should not need this function since the reset function
-    #   one could set the batch size then do reset
-    # @compilable
-    # def batched_reset(self, batch_size):
-    #     in_restVals = jnp.zeros((batch_size, *self.area_shape))      ## input: (B | ix | iy)
-
-    #     out_restVals = jnp.zeros((batch_size,      ## output.shape: (B | n_cells * px * py)
-    #                               self.n_cells * self.patch_shape[0] * self.patch_shape[1]))
-
-    #     self.inputs.set(in_restVals)
-    #     self.outputs.set(out_restVals)
-
     @classmethod
     def help(cls): ## component help function
         properties = {