cleanup, made some fabric routines act on lists of nlm statevectors f…

…or speedup.

README.md

@@ -15,8 +15,10 @@ Spectral CPO model of polycrystalline materials that can:
 
 See [the specfab docs](https://nicholasmr.github.io/specfab) for installation, tutorials, and more.
 
+<!--
 ## Q&A
 
 | **Q** | **A** |
 | :--- | :--- |
 | What $L$ are possible? | Any $4\leq L\leq 20$. If higher $L$ are required: <br>1. `cd src/include && python3 make_gaunt_coefs.py L` (replacing `L`) <br>2. `make clean && make` |
+-->

demo/README.md

@@ -2,7 +2,7 @@
 
 | Folder(s) | Publication |
 | --- | --- |
-| `state-space-validation/olivine`, `enhancement-factor/olivine`| [Rathmann et al. (2024), in prep.|
+| `state-space-validation/olivine`, `enhancement-factor/olivine`| Rathmann et al. (2024), in prep.|
 | `state-space-validation/ice`| [Lilien et al. (2023)](https://doi.org/doi:10.1017/jog.2023.78) |
 | `inverse-elastic-problem` | [Rathmann et al. (2022)](https://doi.org/10.1098/rspa.2022.0574) |
 | `rheology-compare` | [Rathmann and Lilien (2022)](https://doi.org/10.1017/jog.2022.33) |

src/specfabpy/common.py

@@ -15,34 +15,38 @@ def eigenframe(Z, modelplane=None, symframe=-1):
     """
     Principal fabric frame (typically a^(2) principal axes).
     """
-
-    # Is input one-dimensional (vector)? => assume nlm was passed
-    if len(Z.shape) == 1: 
-        Q = sf__.a2(Z) if (symframe == -1) else sf__.a4_eigentensors(Z)[symframe]
+
+    if len(Z.shape) >= 1:
+        # if Z is multidimensional then assume nlm[nn,coefs] was passed
+        nlm = np.array(Z, ndmin=2) # ensure nlm[nn,coef] format even if nlm[coef] was passed 
+        Q = np.array([sf__.a2(nlm[nn,:]) if (symframe == -1) else sf__.a4_eigentensors(nlm[nn,:])[symframe] for nn in range(nlm.shape[0])])
     else:                 
+        Z = np.array(Z, ndmin=2)
         Q = Z # assume Z = a^(2) passed directly
 
-    eigvals, eigvecs = np.linalg.eig(Q)
-    
+    lami, ei = np.linalg.eig(Q)
+
     if modelplane is None: 
-        I = np.flip(eigvals.argsort()) # largest eigenvalue pair is first entry
-
+        I = np.flip(lami.argsort()) # largest eigenvalue pair is first entry
+        ei, lami = ei[:,:,I], lami[:,I]
     else:
-        if   modelplane=='xy': I = abs(eigvecs[2,:]).argsort() # last entry is eigenvector with largest z-value (component out of model plane)
-        elif modelplane=='xz': I = abs(eigvecs[1,:]).argsort() # ... same but for y-value
+        kk = 1 if modelplane=='xz' else 2 # xz else xy
+        for nn in range(lami.shape[0]):
+            lami_, ei_ = lami[nn,:], ei[nn,:,:] # node nn
+            I = abs(ei_[kk,:]).argsort() # last entry is ei with largest y-value (component out of xz modelplane), or z value for xy modelplane
+            II = np.flip(lami_[I[0:2]].argsort()) # largest eigenvalue sorting of *in-plane ei vectors*
+            I[0:2] = I[II] # rearrange in-plane eigenpairs so that largest eigenvalue pair is first
+            lami[nn,:], ei[nn,:,:] = lami_[I], ei_[:,I]
 
-        if eigvals[I[0]] < eigvals[I[1]]: 
-            I[[0,1]] = I[[1,0]] # swap sorted indices so largest eigenvalue entry is first
-
-    return eigvecs[:,I], eigvals[I]
+    return (ei[0], lami[0]) if Z.shape[0]==1 else (ei, lami) # [node], [, xyz comp], eigenpair
 
 
 def pfJ(nlm): 
 
     """
     Pole figure J (pfJ) index
     """
-
+    # @TODO: nlm index ordering should be transposed to match [node, comp] ordering of other routines
     k = np.sqrt(4*np.pi)**2
     if len(nlm.shape) == 2: J_i = [ k*np.sum(np.square(np.absolute(nlm[:,ii]))) for ii in range(nlm.shape[1]) ]
     else:                   J_i =   k*np.sum(np.square(np.absolute(nlm[:,ii])))

src/specfabpy/firedrake/ice.py

@@ -9,8 +9,8 @@
 ------
 The 2D ice-flow model plane is assumed to be xz. 
 This has the benifit of reducing the DOFs of the fabric evolution problem to involve only real numbers (real-valued state vector, s), considerably speeding up the solver.
-If your problem is in fact xy, nothing changes in the way this class is used; the ODFs etc will simply reflect to assumption that flow is in xz rather than xy.
-When visualizing ODFs/fabric quantities, you might therefore want to rotate the frame-of-reference from xz to xy.
+If your problem is in fact xy, nothing changes in the way this class is used; the (M)ODFs etc will simply reflect to assumption that flow is in xz rather than xy.
+When visualizing (M)ODFs/fabric quantities, you might therefore want to rotate the frame-of-reference from xz to xy.
 
 -----------
 DEFINITIONS
@@ -27,7 +27,15 @@
 Lambda0         : CDRX rate factor
 """
 
+"""
+TODO:
+- SSA fabric source/sinks and topo advection terms need to be included 
+- DDRX solver fails (complains nonliner solver diverged, but problem is linearized and works in fenics??)
+- Steady state solver does not work unless RHS is not an empty form. Adding some dummy zero term seems to work..
+"""
+
 import numpy as np
+#import code # code.interact(local=locals())
 from ..specfabpy import specfabpy as sf__
 from .. import common as sfcom
 from firedrake import *
@@ -73,10 +81,10 @@ def __init__(
         self.s  = Function(self.S)
         self.s0 = Function(self.S)
         # Dynamical matrices rows, e.g. M_LROT[i,:] (this account for the real-real coefficient interactions, sufficient for xz problems)
-        self.Mrr_LROT     = [Function(self.S) for ii in np.arange(self.rnlm_len)] # list of M matrix rows
-        self.Mrr_DDRX_src = [Function(self.S) for ii in np.arange(self.rnlm_len)]
-        self.Mrr_CDRX     = [Function(self.S) for ii in np.arange(self.rnlm_len)] 
-        self.Mrr_REG      = [Function(self.S) for ii in np.arange(self.rnlm_len)] 
+        self.Mrr_LROT     = [Function(self.S) for ii in range(self.rnlm_len)] # list of M matrix rows
+        self.Mrr_DDRX_src = [Function(self.S) for ii in range(self.rnlm_len)]
+        self.Mrr_CDRX     = [Function(self.S) for ii in range(self.rnlm_len)] 
+        self.Mrr_REG      = [Function(self.S) for ii in range(self.rnlm_len)] 
 
         ### Derived quantities: viscous anisotropy, a2 eigenvalues, J index, ...
 
@@ -96,12 +104,20 @@ def __init__(
 
         self.rnlm_iso  = [1/np.sqrt(4*np.pi)] + [0]*(self.rnlm_len-1) # Isotropic and normalized state
         self.nlm_dummy = np.zeros((self.nlm_len))
-        self.M_zero    = np.zeros((self.rnlm_len,self.rnlm_len))
+        self.M_zero    = np.zeros((self.numdofs, self.rnlm_len, self.rnlm_len))
+
+        ### Aux/shortcuts
+
+        self.dofs  = np.arange(self.numdofs)
+        self.dofs0 = np.arange(self.numdofs0)
+        self.srng  = np.arange(self.nlm_len)
+        self.srrng = np.arange(self.rnlm_len)
 
         ### Finish up
 
         self.initialize()
-
+
+
     def initialize(self, s0=None):
         s0_ = self.rnlm_iso if s0 is None else s0
         self.s.assign(project(Constant(s0_), self.S))
@@ -123,19 +139,21 @@ def get_nlm(self, *p, xz2xy=False):
 
     def _nlm_nodal(self):
         sp = project(self.s, self.Sd)
-        self.rnlm = np.array([sp.sub(ii).vector()[:] + 0j for ii in range(self.rnlm_len)]) # reduced form (rnlm) per node
-        self.nlm  = np.array([self.sf.rnlm_to_nlm(self.rnlm[:,nn], self.nlm_len) for nn in range(self.numdofs0)]) # full form (nlm) per node (nlm[node,coef])
+        self.rnlm = np.array([sp.sub(ii).vector()[:] + 0j for ii in self.srrng]) # reduced form (rnlm) per node
+        self.nlm  = np.array([self.sf.rnlm_to_nlm(self.rnlm[:,nn], self.nlm_len) for nn in self.dofs0]) # full form (nlm) per node (nlm[node,coef])
 
-    def evolve(self, u, S, dt, iota=+1, Gamma0=None, Lambda0=None, steadystate=False):
+    def evolve(self, *args, **kwargs):
         """
-        Full-Stokes fabric solver
+        The fabric solver, called to step fabric field forward in time by amount dt
+        
+        @TODO: this routine and _get_weakform() can surely be optimized by better choice of linear solver, preconditioning, and not assembling the weak form on every solve call
         """
         self.s0.assign(self.s)
-        F = self._get_weakform(u, S, dt, iota, Gamma0, Lambda0, steadystate=steadystate)
-        solve(lhs(F)==rhs(F), self.s, self.bcs, solver_parameters={'linear_solver':'gmres',}) # Non-symmetric system. Fastest tested are: gmres, bicgstab, tfqmr
+        F = self._get_weakform(*args, **kwargs)
+        solve(lhs(F)==rhs(F), self.s, self.bcs, solver_parameters={'linear_solver':'gmres',}) # non-symmetric system (fastest tested are: gmres, bicgstab, tfqmr)
         self._nlm_nodal()
 
-    def _get_weakform(self, u, S, dt, iota, Gamma0, Lambda0, zeta=0, steadystate=False):
+    def _get_weakform(self, u, S, dt, iota=None, Gamma0=None, Lambda0=None, zeta=0, steadystate=False):
 
         # Crystal processes to include
         ENABLE_LROT = iota is not None
@@ -149,21 +167,20 @@ def _get_weakform(self, u, S, dt, iota, Gamma0, Lambda0, zeta=0, steadystate=Fal
         Sf = project(S, self.G).vector()[:] # deviatoric stress
 
         # Same but in 3D for fabric problem
-        D3f = np.array([self.mat3d(Df[nn]) for nn in np.arange(self.numdofs)])
-        W3f = np.array([self.mat3d(Wf[nn]) for nn in np.arange(self.numdofs)])
-        S3f = np.array([self.mat3d(Sf[nn]) for nn in np.arange(self.numdofs)])
-
-        # Dynamical matrices M_* at each node as np arrays
-        M_LROT     = np.array([self._M_reduced(self.sf.M_LROT,     self.nlm_dummy, D3f[nn], W3f[nn], iota, zeta) for nn in np.arange(self.numdofs)] ) if ENABLE_LROT else self.M_zero
-        M_DDRX_src = np.array([self._M_reduced(self.sf.M_DDRX_src, self.nlm_dummy, S3f[nn])                      for nn in np.arange(self.numdofs)] ) if ENABLE_DDRX else self.M_zero
-        M_REG      = np.array([self._M_reduced(self.sf.M_REG,      self.nlm_dummy, D3f[nn])                      for nn in np.arange(self.numdofs)] ) if ENABLE_REG  else self.M_zero
-
-        # Populate entries of dynamical matrices row-wise
-        kk = 0 # real-real interactions only in xz model frame
-        for ii in np.arange(self.rnlm_len):
-            if ENABLE_LROT: self.Mrr_LROT[ii].vector()[:]     = M_LROT[:,kk,ii,:]
-            if ENABLE_DDRX: self.Mrr_DDRX_src[ii].vector()[:] = M_DDRX_src[:,kk,ii,:]
-            if ENABLE_REG:  self.Mrr_REG[ii].vector()[:]      = M_REG[:,kk,ii,:]
+        D3f = np.array([self.mat3d(Df[nn]) for nn in self.dofs])
+        W3f = np.array([self.mat3d(Wf[nn]) for nn in self.dofs])
+        S3f = np.array([self.mat3d(Sf[nn]) for nn in self.dofs])
+
+        # Dynamical matrices M_* at each node as np arrays (indexing is M[node,row,column])
+        M_LROT     = np.array([self._M_reduced(self.sf.M_LROT,     self.nlm_dummy, D3f[nn], W3f[nn], iota, zeta) for nn in self.dofs] ) if ENABLE_LROT else self.M_zero
+        M_DDRX_src = np.array([self._M_reduced(self.sf.M_DDRX_src, self.nlm_dummy, S3f[nn])                      for nn in self.dofs] ) if ENABLE_DDRX else self.M_zero
+        M_REG      = np.array([self._M_reduced(self.sf.M_REG,      self.nlm_dummy, D3f[nn])                      for nn in self.dofs] ) if ENABLE_REG  else self.M_zero
+
+        # Populate entries of dynamical matrices *row-wise* (row index is ii)
+        for ii in self.srrng:
+            if ENABLE_LROT: self.Mrr_LROT[ii].vector()[:]     = M_LROT[:,ii,:] # all nodes, row=ii, all columns
+            if ENABLE_DDRX: self.Mrr_DDRX_src[ii].vector()[:] = M_DDRX_src[:,ii,:]
+            if ENABLE_REG:  self.Mrr_REG[ii].vector()[:]      = M_REG[:,ii,:]
 
         ### Construct weak form
 
@@ -174,50 +191,47 @@ def _get_weakform(self, u, S, dt, iota, Gamma0, Lambda0, zeta=0, steadystate=Fal
         dtinv = Constant(1/dt)
         if not steadystate:
             F += dtinv * dot( (self.p-self.s0), self.q)*dx
+
+#        F += dot(self.s_null,self.q)*dx # dummy zero term to make rhs(F) work when solving steady-state problem 
 
         # Real space stabilization (Laplacian diffusion)
         if ENABLE_REG:
             F += self.nu_realspace * inner(grad(self.p), grad(self.q))*dx
 
         # Lattice rotation
         if ENABLE_LROT:
-            F += -sum([ dot(self.Mrr_LROT[ii], self.p)*self.qs[ii]*dx for ii in np.arange(self.rnlm_len)])
+            F += -sum([ dot(self.Mrr_LROT[ii], self.p)*self.qs[ii]*dx for ii in self.srrng])
 
         # DDRX
         if ENABLE_DDRX:
             # @TODO NOT WORKING???
-            F_src  = -Gamma0*sum([ dot(self.Mrr_DDRX_src[ii], self.p)*self.qs[ii]*dx for ii in np.arange(self.rnlm_len)]) 
-            F_sink = -Gamma0*sum([ dot(self.Mrr_DDRX_src[0], self.s0)/self.s0[0]*self.ps[ii] * self.qs[ii]*dx for ii in np.arange(self.rnlm_len)]) # nonlinear sink term is linearized around previous solution (self.s0) following Rathmann and Lilien (2021)
+            F_src  = -Gamma0*sum([ dot(self.Mrr_DDRX_src[ii], self.p)*self.qs[ii]*dx for ii in self.srrng]) 
+            F_sink = -Gamma0*sum([ dot(self.Mrr_DDRX_src[0], self.s0)/self.s0[0]*self.ps[ii] * self.qs[ii]*dx for ii in self.srrng]) # nonlinear sink term is linearized around previous solution (self.s0) following Rathmann and Lilien (2021)
             F += F_src - F_sink
 #            F += F_sink
 
         # Orientation space stabilization (hyper diffusion)
-        F += -self.nu_multiplier * sum([ dot(self.Mrr_REG[ii], self.p)*self.qs[ii]*dx for ii in np.arange(self.rnlm_len)])
+        F += -self.nu_multiplier * sum([ dot(self.Mrr_REG[ii], self.p)*self.qs[ii]*dx for ii in self.srrng])
 
         return F
 
     def _M_reduced(self, Mfun, *args, **kwargs):
-        return self.sf.reduce_M(Mfun(*args, **kwargs), self.rnlm_len) # full form -> reduced form of M_*
+        return self.sf.reduce_M(Mfun(*args, **kwargs), self.rnlm_len)[0] # full form -> reduced form of M_*
 
     def mat3d(self, mat2d):
-        return sfcom.mat3d(mat2d, self.modelplane, reshape=True)
+        return sfcom.mat3d(mat2d, self.modelplane, reshape=True) # 2D to 3D strain-rate/stress tensor assuming (i) tr=0 and (ii) out-of-modelplane shear components vanish
 
-    def eigenframe(self, *p, extrapolate=False):
+    def eigenframe(self, *p, **kwargs):
         """
-        Extract a2 eigenvectors (mi) and eigenvalues (lami) at specified point(s)
-        
-        Returns: (mi, lami)
-
+        Extract a2 eigenvectors (mi) and eigenvalues (lami) at specified list of points p=([x1,x2,...], [y1,y2,...])
         Note that you can back-construct a2 = lam1*np.outer(m1,m1) + lam2*np.outer(m2,m2) + lam3*np.outer(m3,m3) 
+                
+        Returns: (mi, lami)
         """
-        xf, yf = np.array(p[0], ndmin=1), np.array(p[1], ndmin=1)
-        if len(xf.shape) > 1:
-            raise ValueError('eigenframe(): please flatten input, only 1D (x,y) is supported')
-        N = len(xf)
-        mi, lami = np.zeros((N,3,3)), np.zeros((N,3))
-        for ii in np.arange(N): 
-            mi[ii], lami[ii] = sfcom.eigenframe(self.get_nlm(xf[ii],yf[ii]), symframe=self.symframe, modelplane=self.modelplane)
-        return (mi[0], lami[0]) if N==1 else (mi, lami)
+        xf, yf = np.array(p[0], ndmin=1), np.array(p[1], ndmin=1) # (point num, x/y value)
+        nlm = np.array([self.get_nlm(xf[pp],yf[pp], **kwargs) for pp in range(len(xf))]) # (point num, nlm coefs)
+        mi, lami = sfcom.eigenframe(nlm, symframe=self.symframe, modelplane=self.modelplane)
+        return (mi, lami)
 
     def pfJ(self, *args, **kwargs):
         """
@@ -235,7 +249,7 @@ def get_E_CAFFE(self, u, Emin=0.1, Emax=10):
         Returns: E_CAFFE
         """
         Df = project(sym(grad(u)), self.Gd).vector()[:]
-        self.E_CAFFE.vector()[:] = np.array([sf__.E_CAFFE(self.mat3d(Df[nn]), self.nlm[nn], Emin, Emax) for nn in np.arange(self.numdofs0)])
+        self.E_CAFFE.vector()[:] = np.array([sf__.E_CAFFE(self.mat3d(Df[nn]), self.nlm[nn], Emin, Emax) for nn in self.dofs0])
         return self.E_CAFFE
 
     def get_Eij(self, Eij_grain, alpha, n_grain, ei=()):   
@@ -247,37 +261,30 @@ def get_Eij(self, Eij_grain, alpha, n_grain, ei=()):
         """
         ### Check args
 
-        if n_grain != 1: 
+        if n_grain != 1:
             raise ValueError('only n_grain = 1 (linear viscous) is supported')
-        if not(0 <= alpha <= 1): 
+        if not(0 <= alpha <= 1):
             raise ValueError('alpha should be between 0 and 1')
 
         ### Calculate Eij etc. using specfabpy per node
 
-        mi   = np.zeros((3, 3, self.numdofs0)) # (xyz component, i-th vector, node)
-        Eij  = np.zeros((6, self.numdofs0))    # (Eij component, node)
-        lami = np.zeros((3, self.numdofs0))    # (i-th a^(2) eigenvalue, node)        
-
-        for nn in np.arange(self.numdofs0): 
-            eigvecs, lami[:,nn] = sfcom.eigenframe(self.nlm[nn,:], symframe=self.symframe, modelplane=self.modelplane)
-            mi[:,0,nn], mi[:,1,nn], mi[:,2,nn] = eigvecs.T if len(ei) == 0 else ei
-            Eij[:,nn] = self.sf.Eij_tranisotropic(self.nlm[nn,:], mi[:,0,nn], mi[:,1,nn], mi[:,2,nn], Eij_grain, alpha, n_grain)
-
-        """
-        The enhancement-factor model depends on effective (homogenized) grain parameters, calibrated against deformation tests.
-        For CPOs far from the calibration states, negative values *may* occur where Eij should tend to zero if truncation L is not large enough.
-        """
+        mi, lami = sfcom.eigenframe(self.nlm, symframe=self.symframe, modelplane=self.modelplane) # (node,xyz,i), (node,i)
+        e1, e2, e3 = mi[:,:,0], mi[:,:,1], mi[:,:,2] if len(ei) == 0 else ei
+        Eij = np.array([self.sf.Eij_tranisotropic(self.nlm[nn,:], e1[nn,:],e2[nn,:],e3[nn,:], Eij_grain, alpha, n_grain) for nn in self.dofs0]) # (node, Voigt-form Eij 1--6)
+
+        # The enhancement-factor model depends on effective (homogenized) grain parameters, calibrated against deformation tests.
+        # For CPOs far from the calibration states, negative values *may* occur where Eij should tend to zero if truncation L is not large enough.
         Eij[Eij<0] = 1e-2 # Set negative E_ij to a very small value (flow inhibiting)
 
         ### Populate functions
 
-        for ii in range(3): # m1,m2,m3
-            self.lami[ii].vector()[:] = lami[ii]
-            for jj in range(3): # x,y,z
-                self.mi[ii].sub(jj).vector()[:] = mi[jj,ii]
+        for ii in range(3): # vec 1,2,3
+            self.lami[ii].vector()[:] = lami[:,ii]
+            for jj in range(3): # comp x,y,z
+                self.mi[ii].sub(jj).vector()[:] = mi[:,jj,ii]
 
         for kk in range(6): 
-            self.Eij[kk].vector()[:] = Eij[kk]
+            self.Eij[kk].vector()[:] = Eij[:,kk]
 
         return (self.mi, self.Eij, self.lami)