<gh_stars>100-1000
#!/usr/bin/env python
# -*- coding: utf-8 -*-
"""ERT modelling operator classes

* base class
* standard BERT modelling class using pygimli.core (C++) functions
* 2.5D non-optimized totalfield forward operator for ERT (reference)
"""

import numpy as np

import pygimli as pg
from pygimli.frameworks import MeshModelling
from .visualization import showERTData

from pygimli import pf


class ERTModellingBase(MeshModelling):
    """Modelling base class for ERT modelling."""

    def __init__(self, **kwargs):
        super(ERTModellingBase, self).__init__(**kwargs)

    def drawData(self, ax, data=None, **kwargs):
        """Draw data in given axe."""
        kwargs['label'] = kwargs.pop('label', pg.unit('res'))
        kwargs['cMap'] = kwargs.pop('cMap', pg.utils.cMap('res'))

        if hasattr(data, '__iter__'):
            vals = data
            data = self.data
        elif data is None:
            data = self.data

        vals = kwargs.pop('vals', data['rhoa'])

        return showERTData(data, vals=vals, ax=ax, **kwargs)

    def drawModel(self, ax, model, **kwargs):
        """Draw the para domain with option model values."""
        kwargs.setdefault('label', pg.unit('res'))
        kwargs.setdefault('cMap', pg.utils.cMap('res'))
        kwargs.setdefault('logScale', True)

        return super().drawModel(ax=ax, model=model, **kwargs)


class ERTModelling(ERTModellingBase):
    """Forward operator for Electrical Resistivity Tomography.

    Note
    ----
    Convention for complex resistiviy inversion:
    We want to use logarithm transformation for the imaginary part of model
    so we need the startmodel to have positive imaginary parts.
    The sign is flipped back to physical correct assumption before we call
    the response function.
    The Jacobian is calculated with negative imaginary parts and will
    be a conjugated complex block matrix for further calulations.
    """

    def __init__(self, sr=True, verbose=False):
        super(ERTModelling, self).__init__()

        # don't use DC*fop or its regionmanager directly
        #
        self._core = None
        if sr:
            self._core = pg.core.DCSRMultiElectrodeModelling(verbose=verbose)
        else:
            self._core = pg.core.DCMultiElectrodeModelling(verbose=verbose)

        self._core.initJacobian()
        self.setJacobian(self._core.jacobian())

        # called from the ERTManager .. needed?
        self.solution = self._core.solution
        self.setComplex = self._core.setComplex
        self.complex = self._core.complex
        self.calculate = self._core.calculate
        self.calcGeometricFactor = self._core.calcGeometricFactor
        self.mapERTModel = self._core.mapERTModel

        self._conjImag = False  # the imaginary parts are flipped for log trans

    def setVerbose(self, v):
        super().setVerbose(v)
        self._core.setVerbose(v)

    def setDefaultBackground(self):
        """Set the default background behaviour."""
        if self.complex():
            self.regionManager().addRegion(3, self._baseMesh, 2)

        regionIds = self.regionManager().regionIdxs()
        pg.info("Found {} regions.".format(len(regionIds)))
        if len(regionIds) > 1:
            bk = pg.sort(regionIds)[0]
            pg.info("Region with smallest marker ({0}) "
                    "set to background".format(bk))
            self.setRegionProperties(bk, background=True)

    def createStartModel(self, dataVals):
        """Create Starting model for ERT inversion."""
        if self.complex():
            dataC = pg.utils.toComplex(dataVals)
            nModel = self.regionManager().parameterCount() // 2
            smRe = np.ones(nModel) * np.median(np.median(dataC.real))
            smIm = np.ones(nModel) * np.median(np.median(dataC.imag))

            if min(smIm) < 0:
                # we want positive phase model
                sm = smRe - 1j * smIm
                pg.info("Model imaginary part being flipped to positive.")
                self._conjImag = True
            else:
                sm = smRe + 1j * smIm

            return pg.utils.squeezeComplex(sm)  # complex impedance
        else:
            return super(ERTModelling, self).createStartModel(dataVals)

    def flipImagPart(self, v):
        """Flip imaginary port (convention)."""
        z = pg.utils.toComplex(v)
        pg.warn('pre min/max={0} / {1} im: {2} / {3}'.format(
            pf(min(z.real)), pf(max(z.real)),
            pf(min(z.imag)), pf(max(z.imag))))

        v = pg.utils.squeezeComplex(pg.utils.toComplex(v),
                                    conj=self._conjImag)

        z = pg.utils.toComplex(v)
        pg.warn('pos min/max={0} / {1} im: {2} / {3}'.format(
            pf(min(z.real)), pf(max(z.real)),
            pf(min(z.imag)), pf(max(z.imag))))
        return v

    def response(self, mod):
        """Forward response (apparent resistivity)."""
        # ensure the mesh is initialized
        self.mesh()
        if self.complex() and self._conjImag:
            pg.warn('flip imaginary part for response calc')
            mod = self.flipImagPart(mod)

        resp = self._core.response(mod)

        if self.complex() and self._conjImag:
            pg.warn('backflip imaginary part after response calc')
            resp = self.flipImagPart(resp)

        return resp

    def createJacobian(self, mod):
        """Compute Jacobian matrix and store but not return."""
        # ensure the mesh is initialized
        self.mesh()
        if self.complex():
            if self._conjImag:
                pg.warn("Flipping imaginary part for jacobian calc")
                mod = self.flipImagPart(mod)

            self._core.createJacobian(mod)
            self._J = pg.utils.squeezeComplex(self._core.jacobian(),
                                              conj=self._conjImag
                                              )
            self.setJacobian(self._J)
            # pg._r("create Jacobian", self, self._J)
            return self._J

        return self._core.createJacobian(mod)

    def setDataPost(self, data):
        """"""
        self._core.setData(data)

    def setMeshPost(self, mesh):
        """"""
        self._core.setMesh(mesh, ignoreRegionManager=True)


class ERTModellingReference(ERTModellingBase):
    """Reference implementation for 2.5D Electrical Resistivity Tomography."""

    def __init__(self, **kwargs):
        super(ERTModelling, self).__init__()

        self.subPotentials = None
        self.lastResponse = None

        # only for mixed boundary hack since this need to know resistivies.
        self.resistivity = None

        # abscissa k and weight for 2.5 inverse cos-transform
        self.k = None
        self.w = None

    def response(self, model):
        """Solve forward task and return apparent resistivity for self.mesh."""
        # NOTE TODO can't be MT until mixed boundary condition depends on
        # self.resistivity
        pg.tic()
        if not self.data.allNonZero('k'):
            pg.error('Need valid geometric factors: "k".')
            pg.warn('Fallback "k" values to -sign("rhoa")')
            self.data.set('k', -pg.math.sign(self.data('rhoa')))

        mesh = self.mesh()

        nDof = mesh.nodeCount()
        elecs = self.data.sensorPositions()

        nEle = len(elecs)
        nData = self.data.size()

        self.resistivity = res = self.createMappedModel(model, -1.0)

        if self.verbose:
            print("Calculate response for model:", min(res), max(res))

        rMin = elecs[0].dist(elecs[1]) / 2.0
        rMax = elecs[0].dist(elecs[-1]) * 2.0

        k, w = self.getIntegrationWeights(rMin, rMax)

        self.k = k
        self.w = w

        # pg.show(mesh, res, label='res')
        # pg.wait()

        rhs = self.createRHS(mesh, elecs)

        # store all potential fields
        u = np.zeros((nEle, nDof))
        self.subPotentials = [pg.Matrix(nEle, nDof) for i in range(len(k))]

        for i, ki in enumerate(k):
            uE = pg.solve(mesh, a=1./res, b=-(ki * ki)/res, f=rhs,
                          bc={'Robin': ['*', self.mixedBC]},
                          userData={'sourcePos': elecs, 'k': ki},
                          verbose=False, stats=0, debug=False)
            self.subPotentials[i] = uE
            u += w[i] * uE

        # collect potential matrix,
        # i.e., potential for all electrodes and all injections
        pM = np.zeros((nEle, nEle))

        for i in range(nEle):
            pM[i] = pg.interpolate(mesh, u[i, :], destPos=elecs)

        # collect resistivity values for all 4 pole measurements
        r = np.zeros(nData)

        for i in range(nData):
            iA = int(self.data('a')[i])
            iB = int(self.data('b')[i])
            iM = int(self.data('m')[i])
            iN = int(self.data('n')[i])

            uAB = pM[iA] - pM[iB]
            r[i] = uAB[iM] - uAB[iN]

        self.lastResponse = r * self.data('k')

        if self.verbose:
            print("Resp min/max: {0} {1} {2}s".format(min(self.lastResponse),
                                                      max(self.lastResponse),
                                                      pg.dur()))

        return self.lastResponse

    def createJacobian(self, model):
        """TODO WRITEME."""
        if self.subPotentials is None:
            self.response(model)

        J = self.jacobian()
        J.resize(self.data.size(), self.regionManager().parameterCount())

        cells = self.mesh().findCellByMarker(0, -1)
        Si = pg.matrix.ElementMatrix()
        St = pg.matrix.ElementMatrix()

        u = self.subPotentials

        pg.tic()
        if self.verbose:
            print("Calculate sensitivity matrix for model: ",
                  min(model), max(model))

        Jt = pg.Matrix(self.data.size(),
                       self.regionManager().parameterCount())

        for kIdx, w in enumerate(self.w):
            k = self.k[kIdx]
            w = self.w[kIdx]

            Jt *= 0.
            A = pg.matrix.ElementMatrixMap()

            for i, c in enumerate(cells):
                modelIdx = c.marker()

                # 2.5D
                Si.u2(c)
                Si *= k * k
                Si += St.ux2uy2uz2(c)

                # 3D
                # Si.ux2uy2uz2(c); w = w* 2

                A.add(modelIdx, Si)

            for dataIdx in range(self.data.size()):

                a = int(self.data('a')[dataIdx])
                b = int(self.data('b')[dataIdx])
                m = int(self.data('m')[dataIdx])
                n = int(self.data('n')[dataIdx])
                Jt[dataIdx] = A.mult(u[kIdx][a] - u[kIdx][b],
                                     u[kIdx][m] - u[kIdx][n])

            J += w * Jt

        m2 = model*model
        k = self.data('k')

        for i in range(J.rows()):
            J[i] /= (m2 / k[i])

        if self.verbose:
            sumsens = np.zeros(J.rows())
            for i in range(J.rows()):
                sumsens[i] = pg.sum(J[i])
            print("sens sum: median = ", pg.math.median(sumsens),
                  " min = ", pg.min(sumsens),
                  " max = ", pg.max(sumsens))

    def calcGeometricFactor(self, data):
        """Calculate geometry factors for a given dataset."""
        if pg.y(data.sensorPositions()) == pg.z(data.sensorPositions()):
            k = np.zeros(data.size())
            for i in range(data.size()):
                a = data.sensorPosition(data('a')[i])
                b = data.sensorPosition(data('b')[i])
                m = data.sensorPosition(data('m')[i])
                n = data.sensorPosition(data('n')[i])
                k[i] = 1./(2.*np.pi) * (1./a.dist(m) - 1./a.dist(n) -
                                        1./b.dist(m) + 1./b.dist(n))
            return k
        else:
            raise BaseException("Please use BERT for non-standard "
                                "data sets" + str(data))

    def uAnalytical(self, p, sourcePos, k):
        """
        Calculate analytical potential for homogeneous halfspace.

            For sigma = 1 [S m]
        """
        r1A = (p - sourcePos).abs()
        # Mirror on surface at depth=0
        r2A = (p - pg.RVector3(1.0, -1.0, 1.0) * sourcePos).abs()

        if r1A > 1e-12 and r2A > 1e-12:
            return (pg.math.besselK0(r1A * k) + pg.math.besselK0(r2A * k)) / \
                    (2.0 * np.pi)
        else:
            return 0.

    def getIntegrationWeights(self, rMin, rMax):
        """TODO WRITEME."""
        nGauLegendre = max(int((6.0 * np.log10(rMax / rMin))), 4)
        nGauLaguerre = 4

        k = pg.Vector()
        w = pg.Vector()

        k0 = 1.0 / (2.0 * rMin)
        pg.GaussLegendre(0.0, 1.0, nGauLegendre, k, w)
        kLeg = k0 * k * k
        wLeg = 2.0 * k0 * k * w / np.pi

        pg.GaussLaguerre(nGauLaguerre, k, w)
        kLag = k0 * (k + 1.0)
        wLag = k0 * np.exp(k) * w / np.pi

        return pg.cat(kLeg, kLag), pg.cat(wLeg, wLag)

    def mixedBC(self, boundary, userData):
        """Apply mixed boundary conditions."""
        if boundary.marker() != pg.core.MARKER_BOUND_MIXED:
            return 0

        sourcePos = pg.center(userData['sourcePos'])

        k = userData['k']
        r1 = boundary.center() - sourcePos

        # Mirror on surface at depth=0
        r2 = boundary.center() - pg.RVector3(1.0, -1.0, 1.0) * sourcePos
        r1A = r1.abs()
        r2A = r2.abs()

        rho = 1.

        if self.resistivity is not None:
            rho = self.resistivity[boundary.leftCell().id()]

        n = boundary.norm()

        if r1A > 1e-12 and r2A > 1e-12:
            # see mod-dc-2d example for robin like BC and the negative sign
            if (pg.math.besselK0(r1A * k) + pg.math.besselK0(r2A * k)) > 1e-12:

                return k / rho * (r1.dot(n) / r1A * pg.math.besselK1(r1A * k) +
                                  r2.dot(n) / r2A * pg.math.besselK1(r2A * k))\
                                / (pg.math.besselK0(r1A * k) +
                                   pg.math.besselK0(r2A * k))
            else:
                return 0.
        else:
            return 0.

    def pointSource(self, cell, f, userData):
        r"""
        Define function for the current source term.

        :math:`\delta(x-pos), \int f(x) \delta(x-pos)=f(pos)=N(pos)`
            Right hand side entries will be shape functions(pos)
        """
        i = userData['i']
        sourcePos = userData['sourcePos'][i]

        if cell.shape().isInside(sourcePos):
            f.setVal(cell.N(cell.shape().rst(sourcePos)), cell.ids())

    def createRHS(self, mesh, elecs):
        """Create right-hand-side vector."""
        rhs = np.zeros((len(elecs), mesh.nodeCount()))
        for i, e in enumerate(elecs):
            c = mesh.findCell(e)
            rhs[i][c.ids()] = c.N(c.shape().rst(e))
        return rhs



if __name__ == "__main__":
    pass

#!/usr/bin/env python
# -*- coding: utf-8 -*-
"""ERT modelling operator classes

* base class
* standard BERT modelling class using pygimli.core (C++) functions
* 2.5D non-optimized totalfield forward operator for ERT (reference)
"""

import numpy as np

import pygimli as pg
from pygimli.frameworks import MeshModelling
from .visualization import showERTData

from pygimli import pf


class ERTModellingBase(MeshModelling):
    """Modelling base class for ERT modelling."""

    def __init__(self, **kwargs):
        super(ERTModellingBase, self).__init__(**kwargs)

    def drawData(self, ax, data=None, **kwargs):
        """Draw data in given axe."""
        kwargs['label'] = kwargs.pop('label', pg.unit('res'))
        kwargs['cMap'] = kwargs.pop('cMap', pg.utils.cMap('res'))

        if hasattr(data, '__iter__'):
            vals = data
            data = self.data
        elif data is None:
            data = self.data

        vals = kwargs.pop('vals', data['rhoa'])

        return showERTData(data, vals=vals, ax=ax, **kwargs)

    def drawModel(self, ax, model, **kwargs):
        """Draw the para domain with option model values."""
        kwargs.setdefault('label', pg.unit('res'))
        kwargs.setdefault('cMap', pg.utils.cMap('res'))
        kwargs.setdefault('logScale', True)

        return super().drawModel(ax=ax, model=model, **kwargs)


class ERTModelling(ERTModellingBase):
    """Forward operator for Electrical Resistivity Tomography.

    Note
    ----
    Convention for complex resistiviy inversion:
    We want to use logarithm transformation for the imaginary part of model
    so we need the startmodel to have positive imaginary parts.
    The sign is flipped back to physical correct assumption before we call
    the response function.
    The Jacobian is calculated with negative imaginary parts and will
    be a conjugated complex block matrix for further calulations.
    """

    def __init__(self, sr=True, verbose=False):
        super(ERTModelling, self).__init__()

        # don't use DC*fop or its regionmanager directly
        #
        self._core = None
        if sr:
            self._core = pg.core.DCSRMultiElectrodeModelling(verbose=verbose)
        else:
            self._core = pg.core.DCMultiElectrodeModelling(verbose=verbose)

        self._core.initJacobian()
        self.setJacobian(self._core.jacobian())

        # called from the ERTManager .. needed?
        self.solution = self._core.solution
        self.setComplex = self._core.setComplex
        self.complex = self._core.complex
        self.calculate = self._core.calculate
        self.calcGeometricFactor = self._core.calcGeometricFactor
        self.mapERTModel = self._core.mapERTModel

        self._conjImag = False  # the imaginary parts are flipped for log trans

    def setVerbose(self, v):
        super().setVerbose(v)
        self._core.setVerbose(v)

    def setDefaultBackground(self):
        """Set the default background behaviour."""
        if self.complex():
            self.regionManager().addRegion(3, self._baseMesh, 2)

        regionIds = self.regionManager().regionIdxs()
        pg.info("Found {} regions.".format(len(regionIds)))
        if len(regionIds) > 1:
            bk = pg.sort(regionIds)[0]
            pg.info("Region with smallest marker ({0}) "
                    "set to background".format(bk))
            self.setRegionProperties(bk, background=True)

    def createStartModel(self, dataVals):
        """Create Starting model for ERT inversion."""
        if self.complex():
            dataC = pg.utils.toComplex(dataVals)
            nModel = self.regionManager().parameterCount() // 2
            smRe = np.ones(nModel) * np.median(np.median(dataC.real))
            smIm = np.ones(nModel) * np.median(np.median(dataC.imag))

            if min(smIm) < 0:
                # we want positive phase model
                sm = smRe - 1j * smIm
                pg.info("Model imaginary part being flipped to positive.")
                self._conjImag = True
            else:
                sm = smRe + 1j * smIm

            return pg.utils.squeezeComplex(sm)  # complex impedance
        else:
            return super(ERTModelling, self).createStartModel(dataVals)

    def flipImagPart(self, v):
        """Flip imaginary port (convention)."""
        z = pg.utils.toComplex(v)
        pg.warn('pre min/max={0} / {1} im: {2} / {3}'.format(
            pf(min(z.real)), pf(max(z.real)),
            pf(min(z.imag)), pf(max(z.imag))))

        v = pg.utils.squeezeComplex(pg.utils.toComplex(v),
                                    conj=self._conjImag)

        z = pg.utils.toComplex(v)
        pg.warn('pos min/max={0} / {1} im: {2} / {3}'.format(
            pf(min(z.real)), pf(max(z.real)),
            pf(min(z.imag)), pf(max(z.imag))))
        return v

    def response(self, mod):
        """Forward response (apparent resistivity)."""
        # ensure the mesh is initialized
        self.mesh()
        if self.complex() and self._conjImag:
            pg.warn('flip imaginary part for response calc')
            mod = self.flipImagPart(mod)

        resp = self._core.response(mod)

        if self.complex() and self._conjImag:
            pg.warn('backflip imaginary part after response calc')
            resp = self.flipImagPart(resp)

        return resp

    def createJacobian(self, mod):
        """Compute Jacobian matrix and store but not return."""
        # ensure the mesh is initialized
        self.mesh()
        if self.complex():
            if self._conjImag:
                pg.warn("Flipping imaginary part for jacobian calc")
                mod = self.flipImagPart(mod)

            self._core.createJacobian(mod)
            self._J = pg.utils.squeezeComplex(self._core.jacobian(),
                                              conj=self._conjImag
                                              )
            self.setJacobian(self._J)
            # pg._r("create Jacobian", self, self._J)
            return self._J

        return self._core.createJacobian(mod)

    def setDataPost(self, data):
        """"""
        self._core.setData(data)

    def setMeshPost(self, mesh):
        """"""
        self._core.setMesh(mesh, ignoreRegionManager=True)


class ERTModellingReference(ERTModellingBase):
    """Reference implementation for 2.5D Electrical Resistivity Tomography."""

    def __init__(self, **kwargs):
        super(ERTModelling, self).__init__()

        self.subPotentials = None
        self.lastResponse = None

        # only for mixed boundary hack since this need to know resistivies.
        self.resistivity = None

        # abscissa k and weight for 2.5 inverse cos-transform
        self.k = None
        self.w = None

    def response(self, model):
        """Solve forward task and return apparent resistivity for self.mesh."""
        # NOTE TODO can't be MT until mixed boundary condition depends on
        # self.resistivity
        pg.tic()
        if not self.data.allNonZero('k'):
            pg.error('Need valid geometric factors: "k".')
            pg.warn('Fallback "k" values to -sign("rhoa")')
            self.data.set('k', -pg.math.sign(self.data('rhoa')))

        mesh = self.mesh()

        nDof = mesh.nodeCount()
        elecs = self.data.sensorPositions()

        nEle = len(elecs)
        nData = self.data.size()

        self.resistivity = res = self.createMappedModel(model, -1.0)

        if self.verbose:
            print("Calculate response for model:", min(res), max(res))

        rMin = elecs[0].dist(elecs[1]) / 2.0
        rMax = elecs[0].dist(elecs[-1]) * 2.0

        k, w = self.getIntegrationWeights(rMin, rMax)

        self.k = k
        self.w = w

        # pg.show(mesh, res, label='res')
        # pg.wait()

        rhs = self.createRHS(mesh, elecs)

        # store all potential fields
        u = np.zeros((nEle, nDof))
        self.subPotentials = [pg.Matrix(nEle, nDof) for i in range(len(k))]

        for i, ki in enumerate(k):
            uE = pg.solve(mesh, a=1./res, b=-(ki * ki)/res, f=rhs,
                          bc={'Robin': ['*', self.mixedBC]},
                          userData={'sourcePos': elecs, 'k': ki},
                          verbose=False, stats=0, debug=False)
            self.subPotentials[i] = uE
            u += w[i] * uE

        # collect potential matrix,
        # i.e., potential for all electrodes and all injections
        pM = np.zeros((nEle, nEle))

        for i in range(nEle):
            pM[i] = pg.interpolate(mesh, u[i, :], destPos=elecs)

        # collect resistivity values for all 4 pole measurements
        r = np.zeros(nData)

        for i in range(nData):
            iA = int(self.data('a')[i])
            iB = int(self.data('b')[i])
            iM = int(self.data('m')[i])
            iN = int(self.data('n')[i])

            uAB = pM[iA] - pM[iB]
            r[i] = uAB[iM] - uAB[iN]

        self.lastResponse = r * self.data('k')

        if self.verbose:
            print("Resp min/max: {0} {1} {2}s".format(min(self.lastResponse),
                                                      max(self.lastResponse),
                                                      pg.dur()))

        return self.lastResponse

    def createJacobian(self, model):
        """TODO WRITEME."""
        if self.subPotentials is None:
            self.response(model)

        J = self.jacobian()
        J.resize(self.data.size(), self.regionManager().parameterCount())

        cells = self.mesh().findCellByMarker(0, -1)
        Si = pg.matrix.ElementMatrix()
        St = pg.matrix.ElementMatrix()

        u = self.subPotentials

        pg.tic()
        if self.verbose:
            print("Calculate sensitivity matrix for model: ",
                  min(model), max(model))

        Jt = pg.Matrix(self.data.size(),
                       self.regionManager().parameterCount())

        for kIdx, w in enumerate(self.w):
            k = self.k[kIdx]
            w = self.w[kIdx]

            Jt *= 0.
            A = pg.matrix.ElementMatrixMap()

            for i, c in enumerate(cells):
                modelIdx = c.marker()

                # 2.5D
                Si.u2(c)
                Si *= k * k
                Si += St.ux2uy2uz2(c)

                # 3D
                # Si.ux2uy2uz2(c); w = w* 2

                A.add(modelIdx, Si)

            for dataIdx in range(self.data.size()):

                a = int(self.data('a')[dataIdx])
                b = int(self.data('b')[dataIdx])
                m = int(self.data('m')[dataIdx])
                n = int(self.data('n')[dataIdx])
                Jt[dataIdx] = A.mult(u[kIdx][a] - u[kIdx][b],
                                     u[kIdx][m] - u[kIdx][n])

            J += w * Jt

        m2 = model*model
        k = self.data('k')

        for i in range(J.rows()):
            J[i] /= (m2 / k[i])

        if self.verbose:
            sumsens = np.zeros(J.rows())
            for i in range(J.rows()):
                sumsens[i] = pg.sum(J[i])
            print("sens sum: median = ", pg.math.median(sumsens),
                  " min = ", pg.min(sumsens),
                  " max = ", pg.max(sumsens))

    def calcGeometricFactor(self, data):
        """Calculate geometry factors for a given dataset."""
        if pg.y(data.sensorPositions()) == pg.z(data.sensorPositions()):
            k = np.zeros(data.size())
            for i in range(data.size()):
                a = data.sensorPosition(data('a')[i])
                b = data.sensorPosition(data('b')[i])
                m = data.sensorPosition(data('m')[i])
                n = data.sensorPosition(data('n')[i])
                k[i] = 1./(2.*np.pi) * (1./a.dist(m) - 1./a.dist(n) -
                                        1./b.dist(m) + 1./b.dist(n))
            return k
        else:
            raise BaseException("Please use BERT for non-standard "
                                "data sets" + str(data))

    def uAnalytical(self, p, sourcePos, k):
        """
        Calculate analytical potential for homogeneous halfspace.

            For sigma = 1 [S m]
        """
        r1A = (p - sourcePos).abs()
        # Mirror on surface at depth=0
        r2A = (p - pg.RVector3(1.0, -1.0, 1.0) * sourcePos).abs()

        if r1A > 1e-12 and r2A > 1e-12:
            return (pg.math.besselK0(r1A * k) + pg.math.besselK0(r2A * k)) / \
                    (2.0 * np.pi)
        else:
            return 0.

    def getIntegrationWeights(self, rMin, rMax):
        """TODO WRITEME."""
        nGauLegendre = max(int((6.0 * np.log10(rMax / rMin))), 4)
        nGauLaguerre = 4

        k = pg.Vector()
        w = pg.Vector()

        k0 = 1.0 / (2.0 * rMin)
        pg.GaussLegendre(0.0, 1.0, nGauLegendre, k, w)
        kLeg = k0 * k * k
        wLeg = 2.0 * k0 * k * w / np.pi

        pg.GaussLaguerre(nGauLaguerre, k, w)
        kLag = k0 * (k + 1.0)
        wLag = k0 * np.exp(k) * w / np.pi

        return pg.cat(kLeg, kLag), pg.cat(wLeg, wLag)

    def mixedBC(self, boundary, userData):
        """Apply mixed boundary conditions."""
        if boundary.marker() != pg.core.MARKER_BOUND_MIXED:
            return 0

        sourcePos = pg.center(userData['sourcePos'])

        k = userData['k']
        r1 = boundary.center() - sourcePos

        # Mirror on surface at depth=0
        r2 = boundary.center() - pg.RVector3(1.0, -1.0, 1.0) * sourcePos
        r1A = r1.abs()
        r2A = r2.abs()

        rho = 1.

        if self.resistivity is not None:
            rho = self.resistivity[boundary.leftCell().id()]

        n = boundary.norm()

        if r1A > 1e-12 and r2A > 1e-12:
            # see mod-dc-2d example for robin like BC and the negative sign
            if (pg.math.besselK0(r1A * k) + pg.math.besselK0(r2A * k)) > 1e-12:

                return k / rho * (r1.dot(n) / r1A * pg.math.besselK1(r1A * k) +
                                  r2.dot(n) / r2A * pg.math.besselK1(r2A * k))\
                                / (pg.math.besselK0(r1A * k) +
                                   pg.math.besselK0(r2A * k))
            else:
                return 0.
        else:
            return 0.

    def pointSource(self, cell, f, userData):
        r"""
        Define function for the current source term.

        :math:`\delta(x-pos), \int f(x) \delta(x-pos)=f(pos)=N(pos)`
            Right hand side entries will be shape functions(pos)
        """
        i = userData['i']
        sourcePos = userData['sourcePos'][i]

        if cell.shape().isInside(sourcePos):
            f.setVal(cell.N(cell.shape().rst(sourcePos)), cell.ids())

    def createRHS(self, mesh, elecs):
        """Create right-hand-side vector."""
        rhs = np.zeros((len(elecs), mesh.nodeCount()))
        for i, e in enumerate(elecs):
            c = mesh.findCell(e)
            rhs[i][c.ids()] = c.N(c.shape().rst(e))
        return rhs



if __name__ == "__main__":
    pass
#!/usr/bin/env python # -*- coding: utf-8 -*- ERT modelling operator classes  * base class * standard BERT modelling class using pygimli.core (C++) functions * 2.5D non-optimized totalfield forward operator for ERT (reference) Modelling base class for ERT modelling. Draw data in given axe. Draw the para domain with option model values. Forward operator for Electrical Resistivity Tomography.      Note     ----     Convention for complex resistiviy inversion:     We want to use logarithm transformation for the imaginary part of model     so we need the startmodel to have positive imaginary parts.     The sign is flipped back to physical correct assumption before we call     the response function.     The Jacobian is calculated with negative imaginary parts and will     be a conjugated complex block matrix for further calulations. # don't use DC*fop or its regionmanager directly # # called from the ERTManager .. needed? # the imaginary parts are flipped for log trans Set the default background behaviour. Create Starting model for ERT inversion. # we want positive phase model # complex impedance Flip imaginary port (convention). Forward response (apparent resistivity). # ensure the mesh is initialized Compute Jacobian matrix and store but not return. # ensure the mesh is initialized # pg._r("create Jacobian", self, self._J) Reference implementation for 2.5D Electrical Resistivity Tomography. # only for mixed boundary hack since this need to know resistivies. # abscissa k and weight for 2.5 inverse cos-transform Solve forward task and return apparent resistivity for self.mesh. # NOTE TODO can't be MT until mixed boundary condition depends on # self.resistivity # pg.show(mesh, res, label='res') # pg.wait() # store all potential fields # collect potential matrix, # i.e., potential for all electrodes and all injections # collect resistivity values for all 4 pole measurements TODO WRITEME. # 2.5D # 3D # Si.ux2uy2uz2(c); w = w* 2 Calculate geometry factors for a given dataset. Calculate analytical potential for homogeneous halfspace.              For sigma = 1 [S m] # Mirror on surface at depth=0 TODO WRITEME. Apply mixed boundary conditions. # Mirror on surface at depth=0 # see mod-dc-2d example for robin like BC and the negative sign Define function for the current source term.          :math:`\delta(x-pos), \int f(x) \delta(x-pos)=f(pos)=N(pos)`             Right hand side entries will be shape functions(pos) Create right-hand-side vector.