removed old weight_distribution.py, other cleanup/revisions throughout

Author: Alexander Ororbia

docs/images/tutorials/neurocog/gmm_fit.jpg

@@ -90,14 +90,6 @@ ngclearn.utils.surrogate\_fx module
    :undoc-members:
    :show-inheritance:
 
-ngclearn.utils.weight\_distribution module
-------------------------------------------
-
-.. automodule:: ngclearn.utils.weight_distribution
-   :members:
-   :undoc-members:
-   :show-inheritance:
-
 Module contents
 ---------------
 

docs/images/tutorials/neurocog/gmm_samples.jpg

@@ -1,9 +1,9 @@
 # Density Modeling and Analysis 
 
 NGC-Learn offers some support for density modeling/estimation, which can be particularly useful in analyzing how internal properties of neuronal models' self-organized cell populations (e.g., how the distributed representations of a model might cluster into distinct groups/categories) or to draw samples from the underlying generative model implied by a particular neuronal structure (e.g., sampling a trained predictive coding generative model). 
-Particularly, within `ngclearn.utils.density`, one can find implementations of mixture models -- such as a mixture-of-Bernoulli or a mixture-of-Gaussians -- which might be employed to carry out such tasks. In this small lesson, we will demonstrate how to set up a Gaussian mixture model (GMM), fit it to some synthetic latent code data, and plot out the distribution it learns overlaid over the data samples as well as examine the kinds of patterns one may sample from the learnt GMM.
+Particularly, within `ngclearn.utils.density`, one can find implementations of mixture models -- such as mixtures of Bernoullis, Gaussians, and exponentials -- which might be employed to carry out such tasks. In this small lesson, we will demonstrate how to set up a Gaussian mixture model (GMM), fit it to some synthetic latent code data, and plot out the distribution it learns overlaid over the data samples as well as examine the kinds of patterns one may sample from the learnt GMM.
 
-## Setting Up a Gaussian Mixture Model
+## Setting Up a Gaussian Mixture Model (GMM)
 
 Let's say you have a two-dimensional dataset of neural code vectors collected from another model you have simulated -- here, we will artificially synthesize this kind of data in this lesson from an "unobserved" trio of multivariate Gaussians (as was done in the t-SNE tutorial) and pretend that this is a set of collected vector measurements. Furthermore, you decide that, after consideration that your data might follow a multi-modal distribution (and reasonably asssuming that multivariate Gaussians might capture most of the inherent structure/shape), you want to fit a GMM to these codes to later on sample from their underlying multi-modal distribution.
 
@@ -63,44 +63,30 @@ model.fit(X, tol=1e-3, verbose=True) ## set verbose to `False` to silence the fi
 which should print to I/O something akin to: 
 
 ```console
-0: Mean-diff = 1.4143142700195312
-1: Mean-diff = 0.15272194147109985
-2: Mean-diff = 0.1888418346643448
-3: Mean-diff = 0.18062230944633484
-4: Mean-diff = 0.15196363627910614
-5: Mean-diff = 0.1135818138718605
-6: Mean-diff = 0.06951556354761124
-7: Mean-diff = 0.03664496913552284
-8: Mean-diff = 0.026161763817071915
-9: Mean-diff = 0.022674376145005226
-10: Mean-diff = 0.021674498915672302
-11: Mean-diff = 0.02205687016248703
-12: Mean-diff = 0.023379826918244362
-13: Mean-diff = 0.02553001046180725
-14: Mean-diff = 0.028586825355887413
+0: Mean-diff = 1.4147894382476807  log(p(X)) = -1706.0753173828125 nats
+1: Mean-diff = 0.14663299918174744  log(p(X)) = -1386.569091796875 nats
+2: Mean-diff = 0.18331432342529297  log(p(X)) = -1359.6962890625 nats
+3: Mean-diff = 0.17693905532360077  log(p(X)) = -1309.736083984375 nats
+4: Mean-diff = 0.1494818776845932  log(p(X)) = -1250.130615234375 nats
+5: Mean-diff = 0.11344392597675323  log(p(X)) = -1221.0008544921875 nats
+6: Mean-diff = 0.07362686842679977  log(p(X)) = -1204.680419921875 nats
+7: Mean-diff = 0.03828870505094528  log(p(X)) = -1192.706298828125 nats
+8: Mean-diff = 0.025705577805638313  log(p(X)) = -1188.51123046875 nats
+9: Mean-diff = 0.021316207945346832  log(p(X)) = -1187.055908203125 nats
+10: Mean-diff = 0.019372563809156418  log(p(X)) = -1186.157470703125 nats
+11: Mean-diff = 0.018868334591388702  log(p(X)) = -1185.443115234375 nats
 ...
 <shortened for brevity>
 ...
-32: Mean-diff = 0.06849467754364014
-33: Mean-diff = 0.06256962567567825
-34: Mean-diff = 0.05789890140295029
-35: Mean-diff = 0.05557262524962425
-36: Mean-diff = 0.05545869469642639
-37: Mean-diff = 0.056351397186517715
-38: Mean-diff = 0.057266443967819214
-39: Mean-diff = 0.05742649361491203
-40: Mean-diff = 0.05546746402978897
-41: Mean-diff = 0.04826011508703232
-42: Mean-diff = 0.03320707008242607
-43: Mean-diff = 0.016994504258036613
-44: Mean-diff = 0.007737572770565748
-45: Mean-diff = 0.0035514419432729483
-46: Mean-diff = 0.0016557337949052453
-47: Mean-diff = 0.0007792692049406469
-Converged after 48 iterations.
+46: Mean-diff = 0.017377303913235664  log(p(X)) = -1062.2596435546875 nats
+47: Mean-diff = 0.007906327955424786  log(p(X)) = -1060.440185546875 nats
+48: Mean-diff = 0.003615213558077812  log(p(X)) = -1060.09130859375 nats
+49: Mean-diff = 0.0016773870447650552  log(p(X)) = -1060.0233154296875 nats
+50: Mean-diff = 0.0007852672133594751  log(p(X)) = -1060.0093994140625 nats
+Converged after 51 iterations.
 ```
 
-In the above instance, notice that our GMM converged early, reaching a good log likelihood in `48` iterations. We can further calculate our final model's log likelihood over the dataset `X` with the following in-built function:
+In the above instance, notice that our GMM converged early, reaching a good, stable log likelihood in `51` iterations. We can further calculate our final model's log likelihood over the dataset `X` with the following in-built function:
 
 ```python
 # Calculate the GMM log likelihood 
@@ -111,10 +97,10 @@ print(f"log[p(X)] = {logPX} nats")
 which will print out the following:
 
 ```console
-log[p(X)] = -423.30889892578125 nats
+log[p(X)] = -1060.006591796875 nats
 ```
 
-(If you add a log-likelihood measurement before you call `.fit()`, you will see that your original log-likelihood is around `-1046.91 nats`.) 
+(If you add a log-likelihood measurement before you call `.fit()`, you will see that your original log-likelihood is around `-1060.01 nats`.) 
 Now, to visualize if our GMM actually capture the underlying multi-modal distribution of our dataset, we may visualize the final GMM with the following plotting code: 
 
 ```python

docs/source/ngclearn.utils.rst

@@ -2,7 +2,7 @@
 from ngclearn import compilable #from ngcsimlib.parser import compilable
 from ngclearn import Compartment #from ngcsimlib.compartment import Compartment
 from ngcsimlib.logger import info
-import ngclearn.utils.weight_distribution as dist
+from ngclearn.utils.distribution_generator import DistributionGenerator
 from ngclearn.components.synapses.convolution.ngcconv import conv2d
 
 from ngclearn.components.jaxComponent import JaxComponent
@@ -80,7 +80,12 @@ def __init__(
 
         ######################### set up compartments ##########################
         tmp_key, *subkeys = random.split(self.key.get(), 4)
-        weights = dist.initialize_params(subkeys[0], filter_init, shape) ## filter tensor
+        #weights = dist.initialize_params(subkeys[0], filter_init, shape)
+        if self.filter_init is None:
+            info(self.name, "is using default weight initializer!")
+            self.filter_init = DistributionGenerator.uniform(0.025, 0.8)
+        weights = self.filter_init(shape, subkeys[0]) ## filter tensor
+
         self.batch_size = batch_size # 1
         ## Compartment setup and shape computation
         _x = jnp.zeros((self.batch_size, x_size, x_size, n_in_chan))
@@ -91,10 +96,10 @@ def __init__(
         self.outputs = Compartment(jnp.zeros(self.out_shape))
         self.weights = Compartment(weights)
         if self.bias_init is None:
-            info(self.name, "is using default bias value of zero (no bias "
-                            "kernel provided)!")
+            info(self.name, "is using default bias value of zero (no bias kernel provided)!")
         self.biases = Compartment(
-            dist.initialize_params(subkeys[2], bias_init, (1, shape[1])) if bias_init else 0.0
+            #dist.initialize_params(subkeys[2], bias_init, (1, shape[1])) if bias_init else 0.0
+            self.bias_init((1, shape[1]), subkeys[2]) if bias_init else 0.0
         )
 
     @compilable

docs/tutorials/neurocog/density_modeling.md

@@ -2,7 +2,7 @@
 from ngclearn import compilable #from ngcsimlib.parser import compilable
 from ngclearn import Compartment #from ngcsimlib.compartment import Compartment
 from ngcsimlib.logger import info
-import ngclearn.utils.weight_distribution as dist
+from ngclearn.utils.distribution_generator import DistributionGenerator
 from ngclearn.components.synapses.convolution.ngcconv import deconv2d
 
 from ngclearn.components.jaxComponent import JaxComponent
@@ -68,8 +68,12 @@ def __init__(
 
         ######################### set up compartments ##########################
         tmp_key, *subkeys = random.split(self.key.get(), 4)
-        weights = dist.initialize_params(subkeys[0], filter_init,
-                                         shape)  ## filter tensor
+        #weights = dist.initialize_params(subkeys[0], filter_init, shape)
+        if self.filter_init is None:
+            info(self.name, "is using default weight initializer!")
+            self.filter_init = DistributionGenerator.uniform(0.025, 0.8)
+        weights = self.filter_init(shape, subkeys[0]) ## filter tensor
+
         self.batch_size = batch_size # 1
         ## Compartment setup and shape computation
         _x = jnp.zeros((self.batch_size, x_size, x_size, n_in_chan))
@@ -82,9 +86,10 @@ def __init__(
         if self.bias_init is None:
             info(self.name, "is using default bias value of zero (no bias "
                             "kernel provided)!")
-        self.biases = Compartment(dist.initialize_params(subkeys[2], bias_init,
-                                                         (1, shape[1]))
-                                  if bias_init else 0.0)
+        self.biases = Compartment(
+            # dist.initialize_params(subkeys[2], bias_init, (1, shape[1])) if bias_init else 0.0
+            self.bias_init((1, shape[1]), subkeys[2]) if bias_init else 0.0
+        )
 
     @compilable
     def advance_state(self):

ngclearn/components/synapses/convolution/convSynapse.py

@@ -75,8 +75,7 @@ def __init__(
         self.weights = Compartment(weights)
         ## Set up (optional) bias values
         if self.bias_init is None:
-            info(self.name, "is using default bias value of zero (no bias "
-                            "kernel provided)!")
+            info(self.name, "is using default bias value of zero (no bias kernel provided)!")
         self.biases = Compartment(self.bias_init((1, shape[1]), subkeys[2]) if bias_init else 0.0)
         # self.biases = Compartment(initialize_params(subkeys[2], bias_init,
         #                                             (1, shape[1]))

ngclearn/components/synapses/convolution/deconvSynapse.py

@@ -11,9 +11,10 @@
 from ngcsimlib import deprecate_args
 
 @partial(jit, static_argnums=[3, 4, 5, 6, 7, 8, 9])
-def _calc_update(pre, post, W, w_bound, is_nonnegative=True, signVal=1.,
-                 prior_type=None, prior_lmbda=0.,
-                 pre_wght=1., post_wght=1.):
+def _calc_update(
+        pre, post, W, w_bound, is_nonnegative=True, signVal=1., prior_type=None, prior_lmbda=0., pre_wght=1.,
+        post_wght=1.
+):
     """
     Compute a tensor of adjustments to be applied to a synaptic value matrix.
 

ngclearn/components/synapses/denseSynapse.py

@@ -5,21 +5,22 @@
 
 from ngclearn.utils.model_utils import clip, d_clip
 import jax
-import jax.numpy as jnp
-import numpy as np
+#import numpy as np
 
 from ngclearn.components.synapses import DenseSynapse
 from ngclearn.utils import tensorstats
 from ngclearn.utils.model_utils import create_function
 
-def gaussian_logpdf(event, mean, stddev):
+def _gaussian_logpdf(event, mean, stddev):
   scale_sqrd = stddev ** 2
   log_normalizer = jnp.log(2 * jnp.pi * scale_sqrd)
   quadratic = (event - mean)**2 / scale_sqrd
   return - 0.5 * (log_normalizer + quadratic)
 
 
-def _compute_update(dt, inputs, rewards, act_fx, weights, seed, mu_act_fx, dmu_act_fx, mu_out_min, mu_out_max, scalar_stddev):
+def _compute_update(
+        dt, inputs, rewards, act_fx, weights, seed, mu_act_fx, dmu_act_fx, mu_out_min, mu_out_max, scalar_stddev
+):
     learning_stddev_mask = jnp.asarray(scalar_stddev <= 0.0, dtype=jnp.float32)
     # (input_dim, output_dim * 2) => (input_dim, output_dim), (input_dim, output_dim)
     W_mu, W_logstd = jnp.split(weights, 2, axis=-1)
@@ -37,7 +38,7 @@ def _compute_update(dt, inputs, rewards, act_fx, weights, seed, mu_act_fx, dmu_a
     sample = jnp.clip(sample, mu_out_min, mu_out_max)
     outputs = sample # the actual action that we take
     # Compute log probability density of the Gaussian
-    log_prob = gaussian_logpdf(sample, fx_mean, std).sum(-1)
+    log_prob = _gaussian_logpdf(sample, fx_mean, std).sum(-1)
     # Compute objective (negative REINFORCE objective)
     objective = (-log_prob * rewards).mean() * 1e-2
 
@@ -65,7 +66,6 @@ def _compute_update(dt, inputs, rewards, act_fx, weights, seed, mu_act_fx, dmu_a
     return dW, objective, outputs
 
 
-
 class REINFORCESynapse(DenseSynapse):
     """
     A stochastic synapse implementing the REINFORCE algorithm (policy gradient method). This synapse
@@ -122,8 +122,10 @@ def __init__(
     ) -> None:
         # This is because we have weights mu and weight log sigma
         input_dim, output_dim = shape
-        super().__init__(name, (input_dim, output_dim * 2), weight_init, None, resist_scale,
-                         p_conn, batch_size=batch_size, **kwargs)
+        super().__init__(
+            name, (input_dim, output_dim * 2), weight_init, None, resist_scale, p_conn,
+            batch_size=batch_size, **kwargs
+        )
 
         ## Synaptic hyper-parameters
         self.shape = shape ## shape of synaptic efficacy matrix
@@ -150,12 +152,8 @@ def __init__(
         self.learning_mask = Compartment(jnp.zeros(()))
         self.seed = Compartment(jax.random.PRNGKey(seed if seed is not None else 42))
 
-
-    # @transition(output_compartments=["weights", "dWeights", "objective", "outputs", "accumulated_gradients", "step_count", "seed"])
-    # @staticmethod
     @compilable
     def evolve(self, dt):
-
         # Get compartment values
         weights = self.weights.get()
         dWeights = self.dWeights.get()
@@ -173,13 +171,13 @@ def evolve(self, dt):
             dt, inputs, rewards, self.act_fx, weights, sub_seed, self.mu_act_fx, self.dmu_act_fx, self.mu_out_min, self.mu_out_max, self.scalar_stddev
         )
         ## do a gradient ascent update/shift
-        weights = (weights + dWeights * self.eta) * self.learning_mask + weights * (1.0 - self.learning_mask) # update the weights only where learning_mask is 1.0
+        weights = (weights + dWeights * self.eta) * self.learning_mask + weights * (1.0 - self.learning_mask.get()) # update the weights only where learning_mask is 1.0
         ## enforce non-negativity
         eps = 0.0 # 0.01 # 0.001
         weights = jnp.clip(weights, eps, self.w_bound - eps)  # jnp.abs(w_bound))
         step_count += 1
         accumulated_gradients = (step_count - 1) / step_count * accumulated_gradients * self.decay + 1.0 / step_count * dWeights # EMA update of accumulated gradients
-        step_count = step_count * (1 - self.learning_mask) # reset the step count to 0 when we have learned
+        step_count = step_count * (1 - self.learning_mask.get()) # reset the step count to 0 when we have learned
 
         # Set updated compartment values
         self.weights.set(weights)
@@ -190,8 +188,6 @@ def evolve(self, dt):
         self.step_count.set(step_count)
         self.seed.set(main_seed)
 
-    # @transition(output_compartments=["inputs", "outputs", "objective", "rewards", "dWeights", "accumulated_gradients", "step_count", "seed"])
-    # @staticmethod
     @compilable
     def reset(self):
         preVals = jnp.zeros((self.batch_size, self.shape[0]))
@@ -214,7 +210,6 @@ def reset(self):
         self.step_count.set(step_count)
         self.seed.set(seed)
 
-
     @classmethod
     def help(cls): ## component help function
         properties = {

ngclearn/components/synapses/hebbian/hebbianSynapse.py

@@ -13,9 +13,10 @@
 from ngclearn.utils import tensorstats
 
 # @partial(jit, static_argnums=[3, 4, 5, 6, 7, 8, 9])
-def _calc_update(pre, post, W, mask, w_bound, is_nonnegative=True, signVal=1.,
-                 prior_type=None, prior_lmbda=0.,
-                 pre_wght=1., post_wght=1.):
+def _calc_update(
+        pre, post, W, mask, w_bound, is_nonnegative=True, signVal=1., prior_type=None, prior_lmbda=0., pre_wght=1.,
+        post_wght=1.
+):
     """
     Compute a tensor of adjustments to be applied to a synaptic value matrix.
 
@@ -190,12 +191,15 @@ class HebbianPatchedSynapse(PatchedSynapse):
         batch_size: the size of each mini batch
     """
 
-    def __init__(self, name, shape, n_sub_models=1, stride_shape=(0,0), eta=0., weight_init=None, bias_init=None,
-                 block_mask=None, w_bound=1., is_nonnegative=False, prior=(None, 0.), sign_value=1.,
-                 optim_type="sgd", pre_wght=1., post_wght=1., p_conn=1.,
-                 resist_scale=1., batch_size=1, **kwargs):
-        super().__init__(name, shape, n_sub_models, stride_shape, block_mask, weight_init, bias_init, resist_scale,
-                         p_conn, batch_size=batch_size, **kwargs)
+    def __init__(
+            self, name, shape, n_sub_models=1, stride_shape=(0,0), eta=0., weight_init=None, bias_init=None,
+            block_mask=None, w_bound=1., is_nonnegative=False, prior=(None, 0.), sign_value=1., optim_type="sgd",
+            pre_wght=1., post_wght=1., p_conn=1., resist_scale=1., batch_size=1, **kwargs
+    ):
+        super().__init__(
+            name, shape, n_sub_models, stride_shape, block_mask, weight_init, bias_init, resist_scale, p_conn,
+            batch_size=batch_size, **kwargs
+        )
 
         prior_type, prior_lmbda = prior
         self.prior_type = prior_type
@@ -338,23 +342,6 @@ def help(cls): ## component help function
                 "hyperparameters": hyperparams}
         return info
 
-
-
-    def __repr__(self):
-        comps = [varname for varname in dir(self) if isinstance(getattr(self, varname), Compartment)]
-        maxlen = max(len(c) for c in comps) + 5
-        lines = f"[{self.__class__.__name__}] PATH: {self.name}\n"
-        for c in comps:
-            stats = tensorstats(getattr(self, c).get())
-            if stats is not None:
-                line = [f"{k}: {v}" for k, v in stats.items()]
-                line = ", ".join(line)
-            else:
-                line = "None"
-            lines += f"  {f'({c})'.ljust(maxlen)}{line}\n"
-        return lines
-
-
 if __name__ == '__main__':
     from ngcsimlib.context import Context
     with Context("Bar") as bar: