text: # Copyright 2022 Google LLC
  #
  # 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
  #
  #     https://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.
  
  """Code to calculate the elastic modulus tensor for athermal systems.
  
    The elastic modulus tensor describes a material's response to different
    boundary deformations. Specifically, for a small deformation given by a
    symmetric strain tensor e, the change in energy is
  
    .. math::
  
       U / V^0 = U^0/V^0 + s^0_{ij} e_{ji} + (1/2) C_{ijkl} e_{ij} e_{kl} + \ldots{}
    
    where :math:`V^0` is the volume, :math:`U^0` is the initial energy, :math:`s^0` is the residual stress
    tensor of the undeformed system, and :math:`C` is the elastic modulus tensor. :math:`C` is
    a fourth-rank tensor of shape (dimension,dimension,dimension,dimension), with
    the following symmetries.
  
      Minor symmetries:
  
      .. math::
  
         C_{ijkl} = C_{jikl} = C_{ijlk}
  
      Major symmetries:
  
      .. math::
  
         C_{ijkl} = C_{lkij}
  
    The minor symmetries are also reflected in the symmetric nature
    of stress and strain tensors:
  
    .. math::
    
       s_{ij} = s_{ji}
       e_{ij} = e_{ji}
  
    In general, there are 21 independent elastic constants in 3 dimension (6 in 2
    dimensions). While systems with additional symmetries (e.g. isotropic,
    orthotropic, etc.) can be expressed with fewer constants, we do not assume any
    such additional symmetries.
  
    At zero temperature, the response of every particle to a deformation can be
    calculated explicitly to linear order, enabling the exact calculation of
    the elastic modulus tensor without the need of any finite differences or any
    approximations.
  
  """
  
  
  
  
  from functools import partial
  from typing import Dict, Callable, Union
  from absl import logging
  
  import jax
  import jax.numpy as jnp
  import jax.scipy as jsp
  from jax import jit, lax, grad, vmap, jacfwd, jacrev, jvp
  
  from jax_md import quantity
  from jax_md.util import Array
  from jax_md.util import f32
  from jax_md.util import f64
  
  def _get_strain_tensor_list(dim, dtype) -> Array:
    if dim == 2:
      strain_tensors = jnp.array([[[1, 0],[0, 0]],
                                  [[0, 0],[0, 1]],
                                  [[0, 1],[1, 0]],
                                  [[1, 0],[0, 1]],
                                  [[1, 1],[1, 0]],
                                  [[0, 1],[1, 1]]], dtype=dtype)
    elif dim == 3:
      strain_tensors = jnp.array([[[1, 0, 0], [0, 0, 0], [0, 0, 0]],
                                  [[0, 0, 0], [0, 1, 0], [0, 0, 0]],
                                  [[0, 0, 0], [0, 0, 0], [0, 0, 1]],
                                  [[0, 0, 0], [0, 0, 1], [0, 1, 0]],
                                  [[0, 0, 1], [0, 0, 0], [1, 0, 0]],
                                  [[0, 1, 0], [1, 0, 0], [0, 0, 0]],
                                  [[0, 0, 0], [0, 1, 0], [0, 0, 1]],
                                  [[1, 0, 0], [0, 0, 0], [0, 0, 1]],
                                  [[1, 0, 0], [0, 1, 0], [0, 0, 0]],
                                  [[1, 0, 0], [0, 0, 1], [0, 1, 0]],
                                  [[1, 0, 1], [0, 0, 0], [1, 0, 0]],
                                  [[1, 1, 0], [1, 0, 0], [0, 0, 0]],
                                  [[0, 0, 0], [0, 1, 1], [0, 1, 0]],
                                  [[0, 0, 1], [0, 1, 0], [1, 0, 0]],
                                  [[0, 1, 0], [1, 1, 0], [0, 0, 0]],
                                  [[0, 0, 0], [0, 0, 1], [0, 1, 1]],
                                  [[0, 0, 1], [0, 0, 0], [1, 0, 1]],
                                  [[0, 1, 0], [1, 0, 0], [0, 0, 1]],
                                  [[0, 0, 1], [0, 0, 1], [1, 1, 0]],
                                  [[0, 1, 0], [1, 0, 1], [0, 1, 0]],
                                  [[0, 1, 1], [1, 0, 0], [1, 0, 0]]], dtype=dtype)
    else:
      raise AssertionError('not implemented for {} dimensions'.format(dim))
    return strain_tensors
  
  def _convert_responses_to_elastic_constants(response_all: Array) -> Array:
    if response_all.shape[0] == 6:
      cxxxx = response_all[0]
      cyyyy = response_all[1]
      cxyxy = 0.25 * response_all[2]
      cxxyy = 0.5  * (response_all[3] - cxxxx - cyyyy)
      cxxxy = 0.25 * (response_all[4] - 4 * cxyxy - cxxxx)
      cyyxy = 0.25 * (response_all[5] - 4 * cxyxy - cyyyy)
  
      C = jnp.array(
          [[[[cxxxx, cxxxy], [cxxxy, cxxyy]],
            [[cxxxy, cxyxy], [cxyxy, cyyxy]]],
           [[[cxxxy, cxyxy], [cxyxy, cyyxy]],
            [[cxxyy, cyyxy], [cyyxy, cyyyy]]]])
  
    elif response_all.shape[0] == 21:
      cxxxx = response_all[0];
      cyyyy = response_all[1];
      czzzz = response_all[2];
      cyzyz = response_all[3] / 4.;
      cxzxz = response_all[4] / 4.;
      cxyxy = response_all[5] / 4.;
      cyyzz = (response_all[6] - cyyyy - czzzz) / 2.;
      cxxzz = (response_all[7] - cxxxx - czzzz) / 2.;
      cxxyy = (response_all[8] - cxxxx - cyyyy) / 2.;
      cxxyz = (response_all[9] - cxxxx - 4. * cyzyz) / 4.;
      cxxxz = (response_all[10] - cxxxx - 4. * cxzxz) / 4.;
      cxxxy = (response_all[11] - cxxxx - 4. * cxyxy) / 4.;
      cyyyz = (response_all[12] - cyyyy - 4. * cyzyz) / 4.;
      cyyxz = (response_all[13] - cyyyy - 4. * cxzxz) / 4.;
      cyyxy = (response_all[14] - cyyyy - 4. * cxyxy) / 4.;
      czzyz = (response_all[15] - czzzz - 4. * cyzyz) / 4.;
      czzxz = (response_all[16] - czzzz - 4. * cxzxz) / 4.;
      czzxy = (response_all[17] - czzzz - 4. * cxyxy) / 4.;
      cyzxz = (response_all[18] - 4. * cyzyz - 4. * cxzxz) / 8.;
      cyzxy = (response_all[19] - 4. * cyzyz - 4. * cxyxy) / 8.;
      cxzxy = (response_all[20] - 4. * cxzxz - 4. * cxyxy) / 8.;
  
      C = jnp.array(
           [[[[cxxxx, cxxxy, cxxxz],
              [cxxxy, cxxyy, cxxyz],
              [cxxxz, cxxyz, cxxzz]],
             [[cxxxy, cxyxy, cxzxy],
              [cxyxy, cyyxy, cyzxy],
              [cxzxy, cyzxy, czzxy]],
             [[cxxxz, cxzxy, cxzxz],
              [cxzxy, cyyxz, cyzxz],
              [cxzxz, cyzxz, czzxz]]],
            [[[cxxxy, cxyxy, cxzxy],
              [cxyxy, cyyxy, cyzxy],
              [cxzxy, cyzxy, czzxy]],
             [[cxxyy, cyyxy, cyyxz],
              [cyyxy, cyyyy, cyyyz],
              [cyyxz, cyyyz, cyyzz]],
             [[cxxyz, cyzxy, cyzxz],
              [cyzxy, cyyyz, cyzyz],
              [cyzxz, cyzyz, czzyz]]],
            [[[cxxxz, cxzxy, cxzxz],
              [cxzxy, cyyxz, cyzxz],
              [cxzxz, cyzxz, czzxz]],
             [[cxxyz, cyzxy, cyzxz],
              [cyzxy, cyyyz, cyzyz],
              [cyzxz, cyzyz, czzyz]],
             [[cxxzz, czzxy, czzxz],
              [czzxy, cyyzz, czzyz],
              [czzxz, czzyz, czzzz]]]])
    else:
      raise AssertionError('response_all has incorrect shape')
    return C
  
  def athermal_moduli(energy_fn: Callable[..., Array],
                      tether_strength: float=1e-10,
                      gradient_check: Array=None,
                      cg_tol: float=1e-7,
                      check_convergence: bool=False
                      ) -> Callable[..., Array]:
    """ Setup calculation of elastic modulus tensor.
  
    Args:
      energy_fn: A function that computes the energy of the system. This
        function must take as an argument `perturbation` which perturbs the
        box shape. Any energy function constructed using `smap` or in `energy.py`
        with a standard space will satisfy this property.
      tether_strength: scalar. Strength of the "tether" applied to each particle,
        which can be necessary to make the Hessian matrix non-singular. Solving
        for the non-affine response of each particle requires that the Hessian is
        positive definite. However, there can often be zero modes (eigenvectors of
        the Hessian with zero eigenvalue) that do not couple to the boundary, and
        therefore do not affect the elastic constants despite the zero eigenvalue.
        The most common example is the global translational modes. To solve for
        the non-affine response, we consider the "tethered Hessian"
          H + tether_strength * I,
        where I is the identity matrix and tether_strength is a small constant.
      gradient_check: None or scalar. If not None, a check will be performed
        to guarantee that the maximum component of the gradient is less than
        gradient_check. In other words, that
          jnp.amax(jnp.abs(grad(energy_fn)(R, box=box))) < gradient_check == True
        NOTE: JAX currently does not support proper runtime error handling.
        Therefore, if this check fails, the calculation will return an array
        of jnp.nan's. It is the users responsibility, if they want to use this
        check, to then ensure that the returned array is not full of nans.
      cg_tol: scalar. Tolerance used when solving for the non-affine response.
      check_convergence: bool. If true, calculate_EMT will return a boolean
        flag specifying if the cg solve routine converged to the desired
        tolerance. The default is False, but convergence checking is highly
        recommended especially when using 32-bit precision data.
  
    Return: A function to calculate the elastic modulus tensor
  
    """
  
    def calculate_emt(R: Array,
                      box: Array,
                      **kwargs) -> Array:
      """Calculate the elastic modulus tensor.
  
      energy_fn(R) corresponds to the state around which we are expanding
  
      Args:
        R: array of shape (N,dimension) of particle positions. This does not
          generalize to arbitrary dimensions and is only implemented for
            dimension == 2
            dimension == 3
        box: A box specifying the shape of the simulation volume. Used to infer
          the volume of the unit cell.
  
      Return: C or the tuple (C,converged)
        where C is the Elastic modulus tensor as an array of shape (dimension,
        dimension,dimension,dimension) that respects the major and minor
        symmetries, and converged is a boolean flag (see above).
  
      """
      if not (R.shape[-1] == 2 or R.shape[-1] == 3):
        raise AssertionError('Only implemented for 2d and 3d systems.')
  
      if R.dtype is not jnp.dtype('float64'):
        logging.warning('Elastic modulus calculations can sometimes lose '
                        'precision when not using 64-bit precision.')
  
      dim = R.shape[-1]
  
      def setup_energy_fn_general(strain_tensor):
        I = jnp.eye(dim, dtype=R.dtype)
        @jit
        def energy_fn_general(R, gamma):
          perturbation = I + gamma * strain_tensor
          return energy_fn(R, perturbation=perturbation, **kwargs)
        return energy_fn_general
  
      def get_affine_response(strain_tensor):
        energy_fn_general = setup_energy_fn_general(strain_tensor)
        d2U_dRdgamma = jacfwd(jacrev(energy_fn_general,argnums=0),argnums=1)(R,0.)
        d2U_dgamma2  = jacfwd(jacrev(energy_fn_general,argnums=1),argnums=1)(R,0.)
        return d2U_dRdgamma, d2U_dgamma2
  
      strain_tensors = _get_strain_tensor_list(dim, R.dtype)
      d2U_dRdgamma_all,d2U_dgamma2_all = vmap(get_affine_response)(strain_tensors)
  
      #Solve the system of equations.
      energy_fn_Ronly = partial(energy_fn, **kwargs)
      def hvp(f, primals, tangents):
        return jvp(grad(f), primals, tangents)[1]
      def hvp_specific_with_tether(v):
        return hvp(energy_fn_Ronly, (R,), (v,)) + tether_strength * v
  
      non_affine_response_all = jsp.sparse.linalg.cg(
          vmap(hvp_specific_with_tether),
          d2U_dRdgamma_all,
          tol=cg_tol
          )[0]
      #The above line should be functionally equivalent to:
      #H0=hessian(energy_fn)(R, box=box, **kwargs).reshape(R.size,R.size) \
      #    + tether_strength * jnp.identity(R.size)
      #non_affine_response_all = jnp.transpose(jnp.linalg.solve(
      #   H0,
      #   jnp.transpose(d2U_dRdgamma_all))
      #   )
  
      residual = jnp.linalg.norm(vmap(hvp_specific_with_tether)(
          non_affine_response_all) - d2U_dRdgamma_all
        )
      converged = residual / jnp.linalg.norm(d2U_dRdgamma_all) < cg_tol
  
      response_all = d2U_dgamma2_all - jnp.einsum("nij,nij->n",
                                                  d2U_dRdgamma_all,
                                                  non_affine_response_all)
  
      vol_0 = quantity.volume(dim, box)
      response_all = response_all / vol_0
      C = _convert_responses_to_elastic_constants(response_all)
  
      # JAX does not allow proper runtime error handling in jitted function.
      # Instead, if the user requests a gradient check and the check fails,
      # we convert C into jnp.nan's. While this doesn't raise an exception,
      # it at least is very "loud".
      if gradient_check is not None:
        maxgrad = jnp.amax(jnp.abs(grad(energy_fn)(R, **kwargs)))
        C = lax.cond(maxgrad > gradient_check,
                     lambda _: jnp.nan * C,
                     lambda _: C,
                     None)
  
      if check_convergence:
        return C, converged
      else:
        return C
  
    return calculate_emt
  
  
  def _get_mandel_mapping_weight(dim, dtype):
    if dim == 2:
      m_map  = jnp.array([[0,0],[1,1],[0,1]], dtype=jnp.int8)
      weight = jnp.array([1,1,jnp.sqrt(2)], dtype=dtype)
      return m_map, weight
    elif dim == 3:
      m_map  = jnp.array([[0,0],[1,1],[2,2],[1,2],[0,2],[0,1]], dtype=jnp.int8)
      weight = jnp.array([1,1,1,jnp.sqrt(2),jnp.sqrt(2),jnp.sqrt(2)], dtype=dtype)
      return m_map, weight
    else:
      raise AssertionError('dim must be 2 or 3')
  
  def tensor_to_mandel(T: Array) -> Array:
    """ Convert a tensor to Mandel notation.
  
    Mandel notation is a way to represent symmetric second-rank tensors and
    fourth-rank tensors with minor symmetries in a reduced form. Pairs of indices
    are combined as follows:
  
        For tensors of shape (2,2) or (2,2,2,2):
        tensor indices        Mandel indices
        0,0 -------------->   0
        1,1 -------------->   1
        0,1 or 1,0 ------->   2
  
        For tensors of shape (3,3) or (3,3,3,3):
        tensor indices        Mandel indices
        0,0 -------------->   0
        1,1 -------------->   1
        2,2 -------------->   2
        1,2 or 2,1 ------->   3
        0,2 or 2,0 ------->   4
        0,1 or 1,0 ------->   5
  
    If mandel_index(i,j) performs the above index mapping, then the input T and
    output M satisfy
      M[mandel_index(i,j)] = T[i,j] * w
    or
      M[mandel_index(i,j), mandel_index(k,l)] = T[i,j,k,l] * w(i,j) * w(k,l)
    where
      w(i,j) = 1       if i==j
             = sqrt(2) if i!=j
    is a weight that is used to ensure proper contraction rules. Here (and only
    here) we do not assume major symmetries in fourth-rank tensors.
  
    Args:
      T: Array with 4 possible shapes:
        1. T.shape == (2,2)
           Convert a symmetric array of shape (2,2) to an array of shape (3,)
        2. T.shape == (3,3)
           Convert a symmetric array of shape (3,3) to an array of shape (6,)
        3. T.shape == (2,2,2,2)
           Convert a tensor of shape (2,2,2,2) with minor symmetries to an array
           of shape (3,3)
        4. T.shape == (3,3,3,3)
           Convert a tensor of shape (3,3,3,3) with minor symmetries to an array
           of shape (6,6)
  
    Output: Array of shape (3,), (6,), (3,3), or (6,6)
  
    see: https://sbrisard.github.io/janus/mandel.html (accessed 21 April, 2021)
    """
    dim = T.shape[0]
    if not (dim ==2 or dim == 3):
      raise AssertionError('dim must be 2 or 3')
  
    rank = len(T.shape)
    if not (rank == 2 or rank == 4):
      raise AssertionError('T must have rank 2 or 4')
  
    m_map, weight = _get_mandel_mapping_weight(dim, T.dtype)
  
    if rank == 2:
      extract = lambda idx, w: T[idx[0], idx[1]] * w
      M = vmap(extract, in_axes=(0,0))(m_map, weight)
    else:
      extract = lambda idx0, idx1, w0, w1: T[
          idx0[0], idx0[1], idx1[0], idx1[1]
          ] * w0 * w1
      M = vmap(
            vmap(extract,
              in_axes=(0,None,0,None)),
            in_axes=(None,0,None,0))(m_map, m_map, weight, weight)
    return M
  
  def mandel_to_tensor(M: Array) -> Array:
    """ Perform the inverse of M = tensor_to_mandel(T).
  
    Args:
      M: Array of shape (3,), (6,), (3,3), or (6,6)
  
    Output: Array of shape (2,2), (3,3), (2,2,2,2), or (3,3,3,3)
    """
    mandel_dim = M.shape[0]
    if not (mandel_dim == 3 or mandel_dim == 6):
      raise AssertionError('M.shape[0] must be 3 or 6')
  
    rank = len(M.shape)
    if not (rank == 1 or rank == 2):
      raise AssertionError('T must have rank 1 or 2')
  
    def mandel_index(i, j):
      return lax.cond(i==j,
                      lambda ij: ij[0],
                      lambda ij: mandel_dim - ij[0] - ij[1],
                      (i,j))
  
    if mandel_dim == 3:
      dimension = 2
    else:
      dimension = 3
  
    tensor_range = jnp.arange(dimension)
    _, weight = _get_mandel_mapping_weight(dimension, M.dtype)
  
    if rank == 1:
      def extract(i, j):
        idx = mandel_index(i, j)
        return M[idx] / weight[idx]
      T = vmap(
            vmap(extract,
              in_axes=(0,None)),
            in_axes=(None,0))(tensor_range, tensor_range)
    else:
      def extract(i, j, k, l):
        idx0 = mandel_index(i, j)
        idx1 = mandel_index(k, l)
        return M[idx0, idx1] / (weight[idx0] * weight[idx1])
      T = vmap(
            vmap(
              vmap(
                vmap(extract,
                  in_axes=(0,None,None,None)),
                in_axes=(None,0,None,None)),
              in_axes=(None,None,0,None)),
            in_axes=(None,None,None,0))(tensor_range,
                                        tensor_range,
                                        tensor_range,
                                        tensor_range)
    return T
  
  
  
  @partial(jit,static_argnums=(1,))
  def _extract_elements(C, as_dict):
    if C.shape[0] == 2:
      indices = ( [0, 1, 0, 0, 0, 0],
                  [0, 1, 1, 0, 0, 1],
                  [0, 1, 0, 1, 0, 1],
                  [0, 1, 1, 1, 1, 1])
      clist = C[ indices ]
      if as_dict:
        names = ['cxxxx','cyyyy','cxyxy','cxxyy','cxxxy','cyyxy']
        return dict(zip(names, clist))
      else:
        return clist
  
    elif C.shape[0] == 3:
      indices = ( [0, 1, 2, 1, 0, 0, 1, 0, 0, 0, 0, 0, 1, 0, 0, 1, 0, 0, 0, 0, 0],
                  [0, 1, 2, 2, 2, 1, 1, 0, 0, 0, 0, 0, 1, 2, 1, 2, 2, 1, 2, 1, 1],
                  [0, 1, 2, 1, 0, 0, 2, 2, 1, 1, 0, 0, 1, 1, 1, 2, 2, 2, 1, 1, 0],
                  [0, 1, 2, 2, 2, 1, 2, 2, 1, 2, 2, 1, 2, 1, 1, 2, 2, 2, 2, 2, 2])
      clist = C[ indices ]
      if as_dict:
        names = ['cxxxx','cyyyy','czzzz','cyzyz','cxzxz','cxyxy','cyyzz','cxxzz',
                 'cxxyy','cxxyz','cxxxz','cxxxy','cyyyz','cyyxz','cyyxy','czzyz',
                 'czzxz','czzxy','cyzxz','cyzxy','cxzxy']
        return dict(zip(names, clist))
      else:
        return clist
    else:
      raise AssertionError('C has wrong shape')
  
  def extract_elements(C: Array) -> Dict:
    """ Convert an elastic modulus tensor into a list of unique elements.
  
        In 2d, these are:
        cxxxx,cyyyy,cxyxy,cxxyy,cxxxy,cyyxy
  
        In 3d, these are:
        cxxxx,cyyyy,czzzz,cyzyz,cxzxz,cxyxy,cyyzz,cxxzz,cxxyy,cxxyz,cxxxz,cxxxy,
        cyyyz,cyyxz,cyyxy,czzyz,czzxz,czzxy,cyzxz,cyzxy,cxzxy
  
    Args:
      C: A previously calculated elastic modulus tensor represented as an
        array of shape (spatial_dimension,spatial_dimension,spatial_dimension,
        spatial_dimension), where spatial_dimension is either 2 or 3. C must
        satisfy both the major and minor symmetries, but this is not checked.
    Return: a dict of the 6 (21) unique elastic constants in 2 (3) dimensions.
    """
    return _extract_elements(C,True)
  
  def extract_isotropic_moduli(C: Array) -> Dict:
    """ Extract commonly used isotropic constants.
  
    There are a number of important constants used to describe the linear
    elastic behavior of isotropic systems, including the bulk modulus, B,
    the shear modulus, G, the longitudinal modulus, M, the Young's modulus,
    E, and the Poisson's ratio, nu. This convenience function extracts them
    from an elastic modulus tensor C.
  
    Angle averaged quantities: While these quantities are defined for
    isotropic systems, one can still define an "angle-averaged shear
    modulus", for example, that averages over all possible shear
    deformations. This can be useful for systems that are statistically
    isotropic but where a particular realization is slightly anisotropic.
    The precise definitions are as follows:
  
    First, we define the "response", R, to a certain strain tensor e to be
      R = 2 * (U / V^0 - U^0/V^0 - s^0_{ij} e_{ji}) = C_{ijkl} e_{ij} e_{kl}
  
    Bulk modulus, B:
    This is the response to the rotationally invariant strain tensor:
        e = (1/2) * ( 1 0 )     or      e = (1/3) * ( 1 0 0 )
                    ( 0 1 )                         ( 0 1 0 )
                                                    ( 0 0 1 )
  
    Shear modulus, G:
    This is the response to the strain tensor:
        e = (1/2) * ( 0 1 )     or      e = (1/2) * ( 0 1 0 )
                    ( 1 0 )                         ( 1 0 0 )
                                                    ( 0 0 0 )
    averaged over all possible orientations. For perfectly isotropic
    systems, it should be equal to C[0,1,0,1].
  
    Longitudinal modulus, M:
    This is the response to the strain tensor:
        e = ( 1 0 )     or      e = ( 1 0 0 )
            ( 0 0 )                 ( 0 0 0 )
                                    ( 0 0 0 )
    averaged over all possible orientations. For perfectly isotropic
    systems, it should be equal to C[0,0,0,0].
  
    Young's modulus, E:
    This is a measure of tensile stiffness and is calculated a using well-
    known expression in terms of B and G.
  
    Poisson's ratio:
    This is a measure of deformations in directions perpendicular to an
    applied load and is calculated a using well-known expression in terms
    of B and G.
  
    Args:
      C: A previously calculated elastic modulus tensor represented as an
        array of shape (spatial_dimension,spatial_dimension,spatial_dimension,
        spatial_dimension), where spatial_dimension is either 2 or 3. C must
        satisfy both the major and minor symmetries, but this is not checked.
  
    Return: a dictionary containing the elastic constants.
  
    """
    if C.shape[0] == 2:
      cxxxx,cyyyy,cxyxy,cxxyy,cxxxy,cyyxy = _extract_elements(C,False)
      B = (cxxxx + cyyyy + 2. * cxxyy) / 4.
      G = (4. * cxyxy + cxxxx + cyyyy - 2. * cxxyy) / 8.
      M = B + G
      E = 4 * B * G / (B + G)
      nu = (B - G) / (B + G)
  
    elif C.shape[0] == 3:
      cxxxx,cyyyy,czzzz,cyzyz,cxzxz,cxyxy,cyyzz,cxxzz,cxxyy, \
      cxxyz,cxxxz,cxxxy,cyyyz,cyyxz,cyyxy,czzyz,czzxz,czzxy, \
      cyzxz,cyzxy,cxzxy = _extract_elements(C,False)
      B = (cxxxx + 2 * cxxyy + 2 * cxxzz + cyyyy + 2 * cyyzz + czzzz) / 9.
      G = (cxxxx - cxxyy - cxxzz + 3 * cxyxy + 3 * cxzxz + cyyyy - cyyzz
           + 3 * cyzyz + czzzz) / 15.
      M = B + 4 * G / 3
      E = 9 * B * G / (3 * B + G)
      nu = (3 * B - 2 * G) / (2 * (3 * B + G))
    else:
      raise AssertionError('C has incorrect shape')
  
    return {'B':B, 'G':G, 'M':M, 'E':E, 'nu':nu}
  
  
  
  
  
  
  
  