moved state perturbations to a separate function, to enable re-use in other parts of the code

models.py

@@ -5,9 +5,6 @@ from scipy.interpolate import CubicSpline
 
 verbose = True
 
-# Initiate a random number generator with preset seed, to ensure reproducibility
-RNG = np.random.default_rng(seed=240718)
-
 class CBL:
     def __init__(self,settings):
         for k, v in settings.items():
@@ -92,7 +89,9 @@ class CBL:
 
         # If desired, add ensemble perturbations to model state
         if self.nens>1 and self.perturb_ensemble_state:
+            x0 = self.perturb_state(x0)
 
+            """
             # Random or uniform perturbations perturbations
             # differ only by magnitude of random sample
             if self.perturbations_type == "random":
@@ -101,10 +100,10 @@ class CBL:
                 randomsize = 1
             
             if self.perturbations_type == "random" or self.perturbations_type == "uniform":
-                ppt = RNG.standard_normal((randomsize,self.nens))
+                ppt = self.RNG.standard_normal((randomsize,self.nens))
                 if self.is_bwind:
-                    ppu = RNG.standard_normal((self.nens,randomsize))
-                    ppv = RNG.standard_normal((self.nens,randomsize))
+                    ppu = self.RNG.standard_normal((self.nens,randomsize))
+                    ppv = self.RNG.standard_normal((self.nens,randomsize))
 
             # Smooth perturbations are slightly more complicated
             if self.perturbations_type == "smooth":
@@ -117,15 +116,15 @@ class CBL:
                 # Draw random perturbations, then interpolate
                 randomsize = self.perturbations_smooth_number
                 ppt = np.zeros((self.nz,self.nens))+np.nan
-                pert_t = RNG.standard_normal((randomsize,self.nens))
+                pert_t = self.RNG.standard_normal((randomsize,self.nens))
                 for n in range(self.nens):
                     f = CubicSpline(ipert,pert_t[:,n])
                     ppt[:,n] = f(np.arange(self.nz))
                 if self.is_bwind:
                     ppu = np.zeros((self.nz,self.nens))+np.nan
                     ppv = np.zeros((self.nz,self.nens))+np.nan
-                    pert_u = RNG.standard_normal((randomsize,self.nens))
-                    pert_v = RNG.standard_normal((randomsize,self.nens))
+                    pert_u = self.RNG.standard_normal((randomsize,self.nens))
+                    pert_v = self.RNG.standard_normal((randomsize,self.nens))
                     for n in range(self.nens):
                         f = CubicSpline(ipert,pert_u[:,n])
                         ppu[:,n] = f(np.arange(self.nz))
@@ -148,6 +147,7 @@ class CBL:
                     if self.is_bwind:
                         x0[self.nz:self.nz*2,n] += self.perturbations_uv_amplitude*ppu[:,n]
                         x0[self.nz*2:self.nz*3,n] += self.perturbations_uv_amplitude*ppv[:,n]
+            """
 
         # If desired, perturb parameters too
         # (last parameter_number elements of state vector)
@@ -215,11 +215,69 @@ class CBL:
         pp = np.zeros((self.parameter_number,self.nens))
 
         for k in range(-self.parameter_number,0):
-            dum = RNG.uniform(self.parameter_ensemble_min[k], self.parameter_ensemble_max[k], size=self.nens)
+            dum = self.RNG.uniform(self.parameter_ensemble_min[k], self.parameter_ensemble_max[k], size=self.nens)
             pp[k,:] = self.parameter_transform[k](dum , kind='dir')
 
         return pp
 
+    def perturb_state(self,x0):
+        # Random or uniform perturbations perturbations
+        # differ only by magnitude of random sample
+        if self.perturbations_type == "random":
+            randomsize = self.nz
+        elif self.perturbations_type == "uniform":
+            randomsize = 1
+        
+        if self.perturbations_type == "random" or self.perturbations_type == "uniform":
+            ppt = self.RNG.standard_normal((randomsize,self.nens))
+            if self.is_bwind:
+                ppu = self.RNG.standard_normal((self.nens,randomsize))
+                ppv = self.RNG.standard_normal((self.nens,randomsize))
+
+        # Smooth perturbations are slightly more complicated
+        if self.perturbations_type == "smooth":
+
+            # Define positions of perturbations
+            # enforce two at boundaries to avoid extrapolation
+            ipert = np.arange(0,self.perturbations_smooth_number-1)*(self.nz//(self.perturbations_smooth_number))
+            ipert = np.append(ipert,self.nz-1)
+
+            # Draw random perturbations, then interpolate
+            randomsize = self.perturbations_smooth_number
+            ppt = np.zeros((self.nz,self.nens))+np.nan
+            pert_t = self.RNG.standard_normal((randomsize,self.nens))
+            for n in range(self.nens):
+                f = CubicSpline(ipert,pert_t[:,n])
+                ppt[:,n] = f(np.arange(self.nz))
+            if self.is_bwind:
+                ppu = np.zeros((self.nz,self.nens))+np.nan
+                ppv = np.zeros((self.nz,self.nens))+np.nan
+                pert_u = self.RNG.standard_normal((randomsize,self.nens))
+                pert_v = self.RNG.standard_normal((randomsize,self.nens))
+                for n in range(self.nens):
+                    f = CubicSpline(ipert,pert_u[:,n])
+                    ppu[:,n] = f(np.arange(self.nz))
+                    f = CubicSpline(ipert,pert_v[:,n])
+                    ppv[:,n] = f(np.arange(self.nz)) 
+
+        if self.perturbations_symmetric:
+            # if n is odd, the last member remains unperturbed
+            for n in range(self.nens//2):
+                x0[:self.nz,n] += self.perturbations_theta_amplitude*ppt[:,n]
+                x0[:self.nz,self.nens//2+n] -= self.perturbations_theta_amplitude*ppt[:,n]
+                if self.is_bwind:
+                    x0[self.nz:self.nz*2,n] += self.perturbations_uv_amplitude*ppu[:,n]
+                    x0[self.nz:self.nz*2,self.nens//2+n] -= self.perturbations_uv_amplitude*ppu[:,n]
+                    x0[self.nz*2:self.nz*3,n] += self.perturbations_uv_amplitude*ppv[:,n]
+                    x0[self.nz*2:self.nz*3,self.nens//2+n] -= self.perturbations_uv_amplitude*ppv[:,n]
+        else:
+            for n in range(self.nens):
+                x0[:self.nz,n] += self.perturbations_theta_amplitude*ppt[:,n]
+                if self.is_bwind:
+                    x0[self.nz:self.nz*2,n] += self.perturbations_uv_amplitude*ppu[:,n]
+                    x0[self.nz*2:self.nz*3,n] += self.perturbations_uv_amplitude*ppv[:,n]
+        return x0
+
     def cbl_model(self,x0,output_full_history=False,perturbed_parameters=np.nan):
         # Read settings in
         nt = self.nt
@@ -310,7 +368,7 @@ class CBL:
                 H0 = Hmax
             
             # Add random perturbations to the initial value
-            H0 += RNG.normal(scale=H0_perturbation_ampl_init)
+            H0 += self.RNG.normal(scale=H0_perturbation_ampl_init)
 
             # Set the surface momentum flux (ustar)
             ustar = 0
@@ -360,7 +418,7 @@ class CBL:
 
             # Add time-dependent surface perturbations
             # Then compute sensible heat flux and integrate T equation
-            H[0,j]     = H0 + RNG.normal(scale=H0_perturbation_ampl_time)
+            H[0,j]     = H0 + self.RNG.normal(scale=H0_perturbation_ampl_time)
             H[1:-1,j]  = -K[1:-1,j]*( (thetap[1:]-thetap[:-1])*rdz - gammac)
             H[nz,j]    = 2*H[nz-1,j]-H[nz-2,j]
             theta[:,j] = thetap[:]-dt*rdz*(H[1:,j]-H[:-1,j])