<commit_before>import time
import sys
import scipy.sparse as sp
import numpy as np
import dask
import dask.array as da
import multiprocessing
from scipy.constants import mu_0
from ...utils.code_utils import deprecate_class

try:
    from pymatsolver import Pardiso as SimpegSolver
except ImportError:
    from SimPEG import SolverLU as SimpegSolver

from ...utils import mkvc, setKwargs, diagEst
from ..frequency_domain.simulation import BaseFDEMSimulation
from ..utils import omega
from .survey import Survey, Data
from pyMKL import mkl_set_num_threads
from .fields import Fields1DPrimarySecondary, Fields3DPrimarySecondary
import xarray as xr


class BaseNSEMSimulation(BaseFDEMSimulation):
    """
    Base class for all Natural source problems.
    """

    # fieldsPair = BaseNSEMFields

    # def __init__(self, mesh, **kwargs):
    #     super(BaseNSEMSimulation, self).__init__()
    #     BaseFDEMProblem.__init__(self, mesh, **kwargs)
    #     setKwargs(self, **kwargs)
    # # Set the default pairs of the problem
    # surveyPair = Survey
    # dataPair = Data

    # Notes:
    # Use the fields and devs methods from BaseFDEMProblem
    def getJtJdiag(self, m, k):
        if k is None:
            k = int(self.survey.nD / 10)

        def JtJv(v):
            Jv = self.Jvec(m, v)
            return self.Jtvec(m, Jv)

        JtJdiag = diagEst(JtJv, len(m), k=k)
        JtJdiag = JtJdiag / np.max(JtJdiag)
        return JtJdiag

    def Jvec(self, m, v, f=None):
        """
        Function to calculate the data sensitivities dD/dm times a vector.

        :param numpy.ndarray m: conductivity model (nP,)
        :param numpy.ndarray v: vector which we take sensitivity product with (nP,)
        :param SimPEG.EM.NSEM.FieldsNSEM (optional) u: NSEM fields object, if not given it is calculated
        :rtype: numpy.ndarray
        :return: Jv (nData,) Data sensitivities wrt m
        """

        # Calculate the fields if not given as input
        if f is None:
            # f = self.fields(m).compute()
            F = self.fields(m)
        # Set current model
        self.model = m
        # Initiate the Jv list
        Jv = []
        number_of_frequencies = len(self.survey.frequencies)
        number_of_components = len(self.survey.get_sources_by_frequency(self.survey.frequencies[0])[0].receiver_list)
        n_dim = number_of_frequencies * number_of_components
        m_dim = int(self.survey.nD / (number_of_components * number_of_frequencies))

        # Loop all the frequenies
        # F = self.fieldsPair(self)
        for nF, freq in enumerate(self.survey.frequencies):
            # Src = self.survey.get_sources_by_frequency(freq)[0]
            # e_s = da.from_delayed(self.fieldByFrequency(freq, nF), (self.mesh.nE, 2), dtype=complex).compute()
            # F[Src, 'e_pxSolution'] = e_s[:, 0]
            # F[Src, 'e_pySolution'] = e_s[:, 1]
            # Get the system
            for src in self.survey.get_sources_by_frequency(freq):
                # need fDeriv_m = df/du*du/dm + df/dm
                # Construct du/dm, it requires a solve
                u_src = F[src, :]
                dA_dm_v = self.getADeriv(freq, u_src, v)
                dRHS_dm_v = self.getRHSDeriv(freq, v)

                # Calculate du/dm*v
                du_dm_v = self.Ainv[nF] * (-dA_dm_v + dRHS_dm_v)
                # Calculate the projection derivatives
                for rx in src.receiver_list:
                    # Calculate dP/du*du/dm*v
                    Jv.append(da.from_delayed(dask.delayed(rx.evalDeriv)(src, self.mesh, F, mkvc(du_dm_v)), shape=(m_dim,), dtype=float))
            # when running full inversion clearing the fields creates error and inversion crashes
            # self.Ainv[nF].clean()
        # return Jv.flatten('F')
        return da.concatenate(Jv, axis=0).compute()

    def Jtvec(self, m, v, f=None):
        """
        Function to calculate the transpose of the data sensitivities (dD/dm)^T times a vector.

        :param numpy.ndarray m: inversion model (nP,)
        :param numpy.ndarray v: vector which we take adjoint product with (nP,)
        :param SimPEG.EM.NSEM.FieldsNSEM f (optional): NSEM fields object, if not given it is calculated
        :rtype: numpy.ndarray
        :return: Jtv (nP,) Data sensitivities wrt m
        """

        if f is None:
            f = self.fields(m).compute()

        self.model = m

        # Ensure v is a data object.
        if not isinstance(v, Data):
            v = Data(self.survey, v)
        # initiate dask array for Jtv
        Jtv = da.zeros(m.size)

        for nF, freq in enumerate(self.survey.frequencies):

            for src in self.survey.get_sources_by_frequency(freq):
                # u_src needs to have both polarizations
                u_src = f[src, :]
                PTv = np.zeros((self.mesh.nE, 2), dtype=complex)
                for rx in src.receiver_list:
                    # Get the adjoint evalDeriv
                    real_or_imag = rx.component
                    if real_or_imag == 'real':
                        PTv += rx.evalDeriv(src, self.mesh, f, mkvc(v[src, rx]), adjoint=True) # wrt f, need possibility wrt m
                    elif real_or_imag == 'imag':
                        PTv += -rx.evalDeriv(src, self.mesh, f, mkvc(v[src, rx]), adjoint=True) # wrt f, need possibility wrt m
                    else:
                        raise Exception('Must be real or imag')
                dA_duIT = da.from_delayed(dask.delayed(mkvc(self.Ainv[nF] * PTv)),
                                          shape=(self.mesh.nE * 2,), dtype=float)  # Force (nU,) shape

                dA_dmT = dask.delayed(self.getADeriv)(freq, u_src, dA_duIT, adjoint=True)
                dRHS_dmT = dask.delayed(self.getRHSDeriv)(freq, dA_duIT, adjoint=True)
                du_dmT = da.from_delayed(-dA_dmT, shape=(self.model.size,),
                                         dtype=complex)
                # Make du_dmT
                du_dmT += da.from_delayed(dRHS_dmT, shape=(self.model.size,), dtype=complex)
                Jtv += du_dmT.real
                # when running full inversion clearing the fields creates error and inversion crashes
                # self.ATinv[nF].clean()
        return Jtv.compute()


###################################
# 1D problems
###################################


class Simulation1DPrimarySecondary(BaseNSEMSimulation):
    """
    A NSEM problem soving a e formulation and primary/secondary fields decomposion.

    By eliminating the magnetic flux density using

        .. math ::

            \mathbf{b} = \\frac{1}{i \omega}\\left(-\mathbf{C} \mathbf{e} \\right)


    we can write Maxwell's equations as a second order system in \\\(\\\mathbf{e}\\\) only:

    .. math ::
        \\left[ \mathbf{C}^{\\top} \mathbf{M_{\mu^{-1}}^e } \mathbf{C} + i \omega \mathbf{M_{\sigma}^f} \\right] \mathbf{e}_{s} = i \omega \mathbf{M_{\sigma_{s}}^f } \mathbf{e}_{p}

    which we solve for :math:`\\mathbf{e_s}`. The total field :math:`\mathbf{e} = \mathbf{e_p} + \mathbf{e_s}`.

    The primary field is estimated from a background model (commonly half space ).


    """
<commit_msg>first pass getJ function<commit_after>import time
import sys
import scipy.sparse as sp
import numpy as np
import dask
import dask.array as da
import multiprocessing
from scipy.constants import mu_0
from ...utils.code_utils import deprecate_class

try:
    from pymatsolver import Pardiso as SimpegSolver
except ImportError:
    from SimPEG import SolverLU as SimpegSolver

from ...utils import mkvc, setKwargs, diagEst
from ..frequency_domain.simulation import BaseFDEMSimulation
from ..utils import omega
from .survey import Survey, Data
from pyMKL import mkl_set_num_threads
from .fields import Fields1DPrimarySecondary, Fields3DPrimarySecondary
import xarray as xr


class BaseNSEMSimulation(BaseFDEMSimulation):
    """
    Base class for all Natural source problems.
    """

    # fieldsPair = BaseNSEMFields

    # def __init__(self, mesh, **kwargs):
    #     super(BaseNSEMSimulation, self).__init__()
    #     BaseFDEMProblem.__init__(self, mesh, **kwargs)
    #     setKwargs(self, **kwargs)
    # # Set the default pairs of the problem
    # surveyPair = Survey
    # dataPair = Data

    # Notes:
    # Use the fields and devs methods from BaseFDEMProblem
    def getJtJdiag(self, m, k):
        if k is None:
            k = int(self.survey.nD / 10)

        def JtJv(v):
            Jv = self.Jvec(m, v)
            return self.Jtvec(m, Jv)

        JtJdiag = diagEst(JtJv, len(m), k=k)
        JtJdiag = JtJdiag / np.max(JtJdiag)
        return JtJdiag

    def Jvec(self, m, v, f=None):
        """
        Function to calculate the data sensitivities dD/dm times a vector.

        :param numpy.ndarray m: conductivity model (nP,)
        :param numpy.ndarray v: vector which we take sensitivity product with (nP,)
        :param SimPEG.EM.NSEM.FieldsNSEM (optional) u: NSEM fields object, if not given it is calculated
        :rtype: numpy.ndarray
        :return: Jv (nData,) Data sensitivities wrt m
        """

        # Calculate the fields if not given as input
        if f is None:
            # f = self.fields(m).compute()
            F = self.fields(m)
        # Set current model
        self.model = m
        # Initiate the Jv list
        Jv = []
        number_of_frequencies = len(self.survey.frequencies)
        number_of_components = len(self.survey.get_sources_by_frequency(self.survey.frequencies[0])[0].receiver_list)
        n_dim = number_of_frequencies * number_of_components
        m_dim = int(self.survey.nD / (number_of_components * number_of_frequencies))

        # Loop all the frequenies
        # F = self.fieldsPair(self)
        for nF, freq in enumerate(self.survey.frequencies):
            # Src = self.survey.get_sources_by_frequency(freq)[0]
            # e_s = da.from_delayed(self.fieldByFrequency(freq, nF), (self.mesh.nE, 2), dtype=complex).compute()
            # F[Src, 'e_pxSolution'] = e_s[:, 0]
            # F[Src, 'e_pySolution'] = e_s[:, 1]
            # Get the system
            for src in self.survey.get_sources_by_frequency(freq):
                # need fDeriv_m = df/du*du/dm + df/dm
                # Construct du/dm, it requires a solve
                u_src = F[src, :]
                dA_dm_v = self.getADeriv(freq, u_src, v)
                dRHS_dm_v = self.getRHSDeriv(freq, v)

                # Calculate du/dm*v
                du_dm_v = self.Ainv[nF] * (-dA_dm_v + dRHS_dm_v)
                # Calculate the projection derivatives
                for rx in src.receiver_list:
                    # Calculate dP/du*du/dm*v
                    Jv.append(da.from_delayed(dask.delayed(rx.evalDeriv)(src, self.mesh, F, mkvc(du_dm_v)), shape=(m_dim,), dtype=float))
            # when running full inversion clearing the fields creates error and inversion crashes
            # self.Ainv[nF].clean()
        # return Jv.flatten('F')
        return da.concatenate(Jv, axis=0).compute()

    def Jtvec(self, m, v, f=None):
        """
        Function to calculate the transpose of the data sensitivities (dD/dm)^T times a vector.

        :param numpy.ndarray m: inversion model (nP,)
        :param numpy.ndarray v: vector which we take adjoint product with (nP,)
        :param SimPEG.EM.NSEM.FieldsNSEM f (optional): NSEM fields object, if not given it is calculated
        :rtype: numpy.ndarray
        :return: Jtv (nP,) Data sensitivities wrt m
        """

        if f is None:
            f = self.fields(m).compute()

        self.model = m

        # Ensure v is a data object.
        if not isinstance(v, Data):
            v = Data(self.survey, v)
        # initiate dask array for Jtv
        Jtv = da.zeros(m.size)

        for nF, freq in enumerate(self.survey.frequencies):

            for src in self.survey.get_sources_by_frequency(freq):
                # u_src needs to have both polarizations
                u_src = f[src, :]
                PTv = np.zeros((self.mesh.nE, 2), dtype=complex)
                for rx in src.receiver_list:
                    # Get the adjoint evalDeriv
                    real_or_imag = rx.component
                    if real_or_imag == 'real':
                        PTv += rx.evalDeriv(src, self.mesh, f, mkvc(v[src, rx]), adjoint=True) # wrt f, need possibility wrt m
                    elif real_or_imag == 'imag':
                        PTv += -rx.evalDeriv(src, self.mesh, f, mkvc(v[src, rx]), adjoint=True) # wrt f, need possibility wrt m
                    else:
                        raise Exception('Must be real or imag')
                dA_duIT = da.from_delayed(dask.delayed(mkvc(self.Ainv[nF] * PTv)),
                                          shape=(self.mesh.nE * 2,), dtype=float)  # Force (nU,) shape

                dA_dmT = dask.delayed(self.getADeriv)(freq, u_src, dA_duIT, adjoint=True)
                dRHS_dmT = dask.delayed(self.getRHSDeriv)(freq, dA_duIT, adjoint=True)
                du_dmT = da.from_delayed(-dA_dmT, shape=(self.model.size,),
                                         dtype=complex)
                # Make du_dmT
                du_dmT += da.from_delayed(dRHS_dmT, shape=(self.model.size,), dtype=complex)
                Jtv += du_dmT.real
                # when running full inversion clearing the fields creates error and inversion crashes
                # self.ATinv[nF].clean()
        return Jtv.compute()


    def getJ(self, m, f=None):
        """
        Function to calculate the sensitivity matrix.

        :param numpy.ndarray m: inversion model (nP,)
        :param SimPEG.EM.NSEM.FieldsNSEM f (optional): NSEM fields object, if not given it is calculated
        :rtype: numpy.ndarray
        :return: J (nD, nP) Data sensitivities wrt m
        """

        if f is None:
            f = self.fields(m)

        self.model = m

        J = np.empty((self.survey.nD, self.model.size))

        istrt = 0
        for freq in self.survey.frequencies:
            AT = self.getA(freq).T

            ATinv = self.Solver(AT, **self.solverOpts)

            for src in self.survey.get_sources_by_frequency(freq):
                # u_src needs to have both polarizations
                u_src = f[src, :]

                for rx in src.receiver_list:
                    # Get the adjoint evalDeriv

                    # Need to make PT
                    # Ideally rx.evalDeriv will be updated to return a matrix,
                    # but for now just calculate it like so...
                    PT = np.empty((AT.shape[0], rx.nD*2), dtype=complex, order='F')
                    for i in range(rx.nD):
                        v = np.zeros(rx.nD)
                        v[i] = 1.0
                        PTv = rx.evalDeriv(src, self.mesh, f, v, adjoint=True)
                        PT[:, 2*i:2*i+2] = PTv

                    dA_duIT = ATinv * PT
                    dA_duIT = dA_duIT.reshape(-1, rx.nD, order='F')  # shape now nUxnD

                    # getADeriv and getRHSDeriv should be updated to accept and return
                    # matrices, but for now this works.
                    # They should also be update to return the real parts as well.
                    dA_dmT = np.empty((rx.nD, J.shape[1]))
                    dRHS_dmT = np.empty((rx.nD, J.shape[1]))
                    for i in range(rx.nD):
                        dA_dmT[i, :] = self.getADeriv(freq, u_src, dA_duIT[:, i], adjoint=True).real
                        dRHS_dmT[i, :] = self.getRHSDeriv(freq, dA_duIT[:, i], adjoint=True).real
                    # Make du_dmT
                    du_dmT = -dA_dmT + dRHS_dmT
                    # Now just need to put it in the right spot.....
                    real_or_imag = rx.component
                    if real_or_imag == 'real':
                        J_rows = du_dmT
                    elif real_or_imag == 'imag':
                        J_rows = -du_dmT
                    else:
                        raise Exception('Must be real or imag')
                    iend = istrt + rx.nD
                    J[istrt:iend, :] = J_rows
                    istrt = iend
            # Clean the factorization, clear memory.
            ATinv.clean()
        return J


###################################
# 1D problems
###################################


class Simulation1DPrimarySecondary(BaseNSEMSimulation):
    """
    A NSEM problem soving a e formulation and primary/secondary fields decomposion.

    By eliminating the magnetic flux density using

        .. math ::

            \mathbf{b} = \\frac{1}{i \omega}\\left(-\mathbf{C} \mathbf{e} \\right)


    we can write Maxwell's equations as a second order system in \\\(\\\mathbf{e}\\\) only:

    .. math ::
        \\left[ \mathbf{C}^{\\top} \mathbf{M_{\mu^{-1}}^e } \mathbf{C} + i \omega \mathbf{M_{\sigma}^f} \\right] \mathbf{e}_{s} = i \omega \mathbf{M_{\sigma_{s}}^f } \mathbf{e}_{p}

    which we solve for :math:`\\mathbf{e_s}`. The total field :math:`\mathbf{e} = \mathbf{e_p} + \mathbf{e_s}`.

    The primary field is estimated from a background model (commonly half space ).


    """

{'fraction_non_alphanumeric': 0.1122478637446957, 'fraction_numerical': 0.002732081613672034, 'mean_word_length': 1.941860465116279, 'pattern_counts': {'":': 0, '<': 3, '<?xml version=': 0, '>': 3, 'https://': 0, 'lorem ipsum': 0, 'www.': 0, 'xml': 0}, 'pii_count': 0, 'substrings_counts': 0, 'word_list_counts': {'cursed_substrings.json': 9, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
{'dir': 'git-commits-cleaned', 'id': '410518', 'n_tokens_mistral': 5631, 'n_tokens_neox': 5282, 'n_words': 1449}