text: #! /usr/bin/env python3
  #
  #  Copyright 2018 California Institute of Technology
  #
  #  Licensed under the Apache License, Version 2.0 (the "License");
  #  you may not use this file except in compliance with the License.
  #  You may obtain a copy of the License at
  #
  #      http://www.apache.org/licenses/LICENSE-2.0
  #
  #  Unless required by applicable law or agreed to in writing, software
  #  distributed under the License is distributed on an "AS IS" BASIS,
  #  WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
  #  See the License for the specific language governing permissions and
  #  limitations under the License.
  #
  # ISOFIT: Imaging Spectrometer Optimal FITting
  # Author: David R Thompson, david.r.thompson@jpl.nasa.gov
  #
  
  import scipy as s
  from scipy.io import loadmat, savemat
  from common import recursive_replace, svd_inv, eps
  from scipy.linalg import block_diag, det, norm, pinv, sqrtm, inv
  from surf import Surface
  
  
  class MultiComponentSurface(Surface):
      """A model of the surface based on a collection of multivariate 
         Gaussians, with one or more equiprobable components and full 
         covariance matrices. 
  
         To evaluate the probability of a new spectrum, we calculate the
         Mahalanobis distance to each component cluster, and use that as our
         Multivariate Gaussian surface model."""
  
      def __init__(self, config):
  
          Surface.__init__(self, config)
  
          # Models are stored as dictionaries in .mat format
          model_dict = loadmat(config['surface_file'])
          self.components = list(zip(model_dict['means'], model_dict['covs']))
          self.n_comp = len(self.components)
          self.wl = model_dict['wl'][0]
          self.n_wl = len(self.wl)
  
          # Set up normalization method
          self.normalize = model_dict['normalize']
          if self.normalize == 'Euclidean':
              self.norm = lambda r: norm(r)
          elif self.normalize == 'RMS':
              self.norm = lambda r: s.sqrt(s.mean(pow(r, 2)))
          elif self.normalize == 'None':
              self.norm = lambda r: 1.0
          else:
              raise ValueError('Unrecognized Normalization: %s\n' %
                               self.normalize)
  
          try:
              self.selection_metric = config['selection_metric']
          except KeyError:
              self.selection_metric = 'Mahalanobis'
  
          # Reference values are used for normalizing the reflectances.
          # in the VSWIR regime, reflectances are normalized so that the model
          # is agnostic to absolute magnitude.
          self.refwl = s.squeeze(model_dict['refwl'])
          self.idx_ref = [s.argmin(abs(self.wl-w))
                          for w in s.squeeze(self.refwl)]
          self.idx_ref = s.array(self.idx_ref)
  
          # Cache some important computations
          self.Covs, self.Cinvs, self.mus = [], [], []
          for i in range(self.n_comp):
              Cov = self.components[i][1]
              self.Covs.append(s.array([Cov[j, self.idx_ref]
                                        for j in self.idx_ref]))
              self.Cinvs.append(svd_inv(self.Covs[-1]))
              self.mus.append(self.components[i][0][self.idx_ref])
  
          # Variables retrieved: each channel maps to a reflectance model parameter
          rmin, rmax = 0, 10.0
          self.statevec = ['RFL_%04i' % int(w) for w in self.wl]
          self.bounds = [[rmin, rmax] for w in self.wl]
          self.scale = [1.0 for w in self.wl]
          self.init = [0.15 * (rmax-rmin)+rmin for v in self.wl]
          self.idx_lamb = s.arange(self.n_wl)
          self.n_state = len(self.statevec)
  
      def component(self, x_surface, geom):
          """ We pick a surface model component using the Mahalanobis distance.
              This always uses the Lambertian (non-specular) version of the 
              surface reflectance."""
  
          if self.n_comp <= 1:
              return 0
  
          # Get the (possibly normalized) reflectance
          lamb = self.calc_lamb(x_surface, geom)
          lamb_ref = lamb[self.idx_ref]
          lamb_ref = lamb_ref / self.norm(lamb_ref)
  
          # Mahalanobis or Euclidean distances
          mds = []
          for ci in range(self.n_comp):
              ref_mu = self.mus[ci]
              ref_Cinv = self.Cinvs[ci]
              if self.selection_metric == 'Mahalanobis':
                  md = (lamb_ref - ref_mu).T.dot(ref_Cinv).dot(lamb_ref - ref_mu)
              else:
                  md = sum(pow(lamb_ref - ref_mu, 2))
              mds.append(md)
          closest = s.argmin(mds)
          return closest
  
      def xa(self, x_surface, geom):
          '''Mean of prior distribution, calculated at state x.  We find
          the covariance in a normalized space (normalizing by z) and then un-
          normalize the result for the calling function.  This always uses the
          Lambertian (non-specular) version of the surface reflectance.'''
  
          lamb = self.calc_lamb(x_surface, geom)
          lamb_ref = lamb[self.idx_ref]
          mu = s.zeros(self.n_state)
          ci = self.component(x_surface, geom)
          lamb_mu = self.components[ci][0]
          lamb_mu = lamb_mu * self.norm(lamb_ref)
          mu[self.idx_lamb] = lamb_mu
          return mu
  
      def Sa(self, x_surface, geom):
          '''Covariance of prior distribution, calculated at state x.  We find
          the covariance in a normalized space (normalizing by z) and then un-
          normalize the result for the calling function.'''
  
          lamb = self.calc_lamb(x_surface, geom)
          lamb_ref = lamb[self.idx_ref]
          ci = self.component(x_surface, geom)
          Cov = self.components[ci][1]
          Cov = Cov * (self.norm(lamb_ref)**2)
  
          # If there are no other state vector elements, we're done.
          if len(self.statevec) == len(self.idx_lamb):
              return Cov
  
          # Embed into a larger state vector covariance matrix
          nprefix = self.idx_lamb[0]
          nsuffix = len(self.statevec)-self.idx_lamb[-1]-1
          Cov_prefix = s.zeros((nprefix, nprefix))
          Cov_suffix = s.zeros((nsuffix, nsuffix))
          return block_diag(Cov_prefix, Cov, Cov_suffix)
  
      def fit_params(self, rfl_meas, Ls, geom):
          '''Given a reflectance estimate, fit a state vector.'''
  
          x_surface = s.zeros(len(self.statevec))
          if len(rfl_meas) != len(self.idx_lamb):
              raise ValueError('Mismatched reflectances')
          for i, r in zip(self.idx_lamb, rfl_meas):
              x_surface[i] = max(self.bounds[i][0]+0.001,
                                 min(self.bounds[i][1]-0.001, r))
          return x_surface
  
      def calc_rfl(self, x_surface, geom):
          '''Non-Lambertian reflectance'''
  
          return self.calc_lamb(x_surface, geom)
  
      def calc_lamb(self, x_surface, geom):
          '''Lambertian reflectance'''
  
          return x_surface[self.idx_lamb]
  
      def drfl_dsurface(self, x_surface, geom):
          '''Partial derivative of reflectance with respect to state vector, 
          calculated at x_surface.'''
  
          return self.dlamb_dsurface(x_surface, geom)
  
      def dlamb_dsurface(self, x_surface, geom):
          '''Partial derivative of Lambertian reflectance with respect to 
             state vector, calculated at x_surface.'''
          dlamb = s.eye(self.n_wl, dtype=float)
          nprefix = self.idx_lamb[0]
          nsuffix = self.n_state - self.idx_lamb[-1] - 1
          prefix = s.zeros((self.n_wl, nprefix))
          suffix = s.zeros((self.n_wl, nsuffix))
          return s.concatenate((prefix, dlamb, suffix), axis=1)
  
      def calc_Ls(self, x_surface, geom):
          '''Emission of surface, as a radiance'''
          return s.zeros(self.n_wl, dtype=float)
  
      def dLs_dsurface(self, x_surface, geom):
          '''Partial derivative of surface emission with respect to state vector, 
          calculated at x_surface.'''
          dLs = s.zeros((self.n_wl, self.n_wl), dtype=float)
          nprefix = self.idx_lamb[0]
          nsuffix = len(self.statevec)-self.idx_lamb[-1]-1
          prefix = s.zeros((self.n_wl, nprefix))
          suffix = s.zeros((self.n_wl, nsuffix))
          return s.concatenate((prefix, dLs, suffix), axis=1)
  
      def summarize(self, x_surface, geom):
          '''Summary of state vector'''
  
          if len(x_surface) < 1:
              return ''
          return 'Component: %i' % self.component(x_surface, geom)
  