def extract_configs_from_scheme(scheme: pb.DataContainerERT,
                                config_indices: list,
                                tmp_dir: str,
                                remove_tmp_file=True):
    """ Extract specific electrode configurations from scheme.

    Utility function to extract specific electrode configurations (identified by index) from a given scheme. Extracted
    data will be returned in a separate scheme.

    Parameter:
        scheme: The scheme the electrode configurations should be extracted from.
        config_indices: The indices of the configurations to be extracted.
        tmp_dir: Directory to temporarily store files in.
        remove_tmp_file: (optional) Whether the temporary file(s) should be removed when no longer necessary. Should
                         only be disabled for debugging purposes.

    Returns:
        A scheme containing the extracted configurations.
    """
    logging.info('Extracting configs from scheme...')
    tmp_file = tmp_dir + 'tmpExtract.shm'
    f = open(tmp_file, 'w')
    # Write electrodes to file
    electrodes = scheme.sensorPositions()
    f.write('%d\n' % len(electrodes))
    f.write('# x y z\n')
    for i in range(len(electrodes)):
        f.write('%f %f %f\n' %
                (electrodes[i][0], electrodes[i][1], electrodes[i][2]))
    # Extract configurations and write to file
    f.write('%d\n' % len(config_indices))
    f.write('# a b m n err i ip iperr k r rhoa u valid \n')
    for i in config_indices:
        f.write('%d %d %d %d %f %f %f %f %f %f %f %f %d\n' %
                (scheme('a')[i] + 1, scheme('b')[i] + 1, scheme('m')[i] + 1,
                 scheme('n')[i] + 1, scheme('err')[i], scheme('i')[i],
                 scheme('ip')[i], scheme('iperr')[i], scheme('k')[i],
                 scheme('r')[i], scheme('rhoa')[i], scheme('u')[i],
                 scheme('valid')[i]))
    f.write('0\n')
    f.close()
    # Read file again to not generate any inconsistencies
    result = pb.load(tmp_file)
    if remove_tmp_file:
        os.remove(tmp_file)
    return result
def invert_pygimli_ERT(inputfileERT, sensors, mesh, date):
    #%% INVERT ERT data

    invpath = main + '/invdir/pygimli/' + date + '/'
    if not os.path.exists(invpath):
        os.makedirs(invpath)

    dataERT = pb.load(main + './raw_data/' + inputfileERT)
    dataERT.setSensorPositions(sensors)

    ert = pg.physics.ERTManager(
        dataERT)  # sr=False, verbose=True, debug=False)
    k = createGeometricFactors(dataERT)
    dataERT.set("k", k)
    dataERT.set('r', dataERT('u') / dataERT('i'))
    dataERT.set('rhoa', dataERT('r') * dataERT('k'))
    dataERT['err'] = ert.estimateError(dataERT,
                                       absoluteError=0.001,
                                       relativeError=0.1)
    dataERT.markInvalid(dataERT("rhoa") < 0)
    dataERT.removeInvalid()
    dataERT.save(invpath + 'dataERT.data')
    # ert.setMesh(mesh3d)  # important to do it explicitly as it avoids rebuilding region manager etc.
    # C = pg.matrix.GeostatisticConstraintsMatrix(mesh=ert.fop.paraDomain, I=[25, 5], dip=-25)
    # ert.fop.setConstraints(C)

    # Set refernce nodes in corners (necessary for closed geometries)
    lower_left_node = mesh.findNearestNode([mesh.xmin(), mesh.ymin(), 0.0])
    mesh.node(lower_left_node).setMarker(-1000)  #MARKER_NODE_CALIBRATION

    lower_right_node = mesh.findNearestNode([mesh.xmax(), mesh.ymin(), 0.0])
    mesh.node(lower_right_node).setMarker(
        -999)  #MARKER_NODE_REFERENCEELECTRODE

    model = ert.invert(mesh=mesh, lam=20, verbose=True)
    pg.info('Inversion stopped with chi² = {0:.3}'.format(ert.fw.chi2()))

    return model, dataERT
import numpy as np

import pybert as pb
import pygimli as pg
pg.verbose = print  # temporary
import pygimli.meshtools as mt
from settings import depth_5000, depth_5198, erte, rste

ertData = pb.load("ert.data")

print("Number of electrodes:", ertData.sensorCount())
print(ertData)

rstData = pg.DataContainer("rst.data", "s g")
print("Number of shot/receivers:", rstData.sensorCount())
maxrst = pg.max(pg.x(rstData.sensors()))

idx = []
for i, sensor in enumerate(ertData.sensors()):
    if sensor[0] >= 50.0:
        idx.append(i)

ertData.removeSensorIdx(idx)
ertData.removeInvalid()
ertData.removeUnusedSensors()
ertData.set("err", pg.RVector(ertData.size(), erte))
ertData.save("ert_filtered.data")

rstData.set("err", pg.RVector(rstData.size(), rste))
#
# # Remove two data points with high v_a at zero-offset
import pybert as pb
import pygimli as pg
import pygimli.meshtools as mt
from pygimli.physics import ERTManager
from pygimli.physics import TravelTimeManager
from pygimli.physics.traveltime.ratools import createGradientModel2D

############
# Settings
maxIter = 15
zWeight = 0.25
############

case = int(sys.argv[1])
ertData = pb.load("erttrue.dat")
print(ertData)
mesh = pg.load("mesh.bms")
depth = mesh.ymax() - mesh.ymin()

# %% Build inversion meshes
plc = mt.createParaMeshPLC(ertData,
                           paraDepth=depth,
                           paraDX=0.3,
                           boundary=0,
                           paraBoundary=2)

if case == 1:
    for depth in (5, 15):
        start = plc.createNode(mesh.xmin(), -depth, 0.0)
        end = plc.createNode(mesh.xmax(), -depth, 0.0)
def merge_schemes(scheme1: pb.DataContainerERT,
                  scheme2: pb.DataContainerERT,
                  tmp_dir: str,
                  remove_tmp_file=True):
    """ Merges to schemes while prioritizing the first one.

    Utility function to merge to schemes. Electrode positions can differ. Electrodes on same positions will be merged.
    When multiple measurements are available for the same electrode configuration, the data from scheme1 will be used.

    Parameter:
        scheme1: First scheme to be merged. This scheme will be prioritized.
        scheme2: Second scheme to be merged.
        tmp_dir: Directory to temporarily store files in.
        remove_tmp_file: (optional) Whether the temporary file(s) should be removed when no longer necessary. Should
                         only be disabled for debugging purposes.

    Returns:
        The merged scheme.
    """
    logging.info('Merging datasets with %d and %d entries',
                 len(scheme1('rhoa')), len(scheme2('rhoa')))
    # Create global electrode register
    electrodes = []
    scheme1_electrodes = scheme1.sensorPositions()
    for i in range(len(scheme1_electrodes)):
        electrodes.append(scheme1_electrodes[i])
    scheme2_electrodes = scheme2.sensorPositions()
    for i in range(len(scheme2_electrodes)):
        if scheme2_electrodes[i] not in electrodes:
            electrodes.append(scheme2_electrodes[i])
    # Create electrode dictionaries
    scheme1_dic = {-1: -1}
    for i in range(len(scheme1_electrodes)):
        scheme1_dic.update({i: i})
    scheme2_dic = {-1: -1}
    for i in range(len(scheme2_electrodes)):
        for j in range(len(electrodes)):
            if scheme2_electrodes[i] == electrodes[j]:
                scheme2_dic.update({i: j})
    # Merge data
    # (sensors: a b m n)
    # (data: err i ip iperr k r rhoa u valid)
    scheme_sensors = []
    scheme_data = []
    for i in range(len(scheme1('rhoa'))):
        scheme_sensors.append([
            scheme1_dic[scheme1('a')[i]] + 1, scheme1_dic[scheme1('b')[i]] + 1,
            scheme1_dic[scheme1('m')[i]] + 1, scheme1_dic[scheme1('n')[i]] + 1
        ])
        scheme_data.append([
            scheme1('err')[i],
            scheme1('i')[i],
            scheme1('ip')[i],
            scheme1('iperr')[i],
            scheme1('k')[i],
            scheme1('r')[i],
            scheme1('rhoa')[i],
            scheme1('u')[i],
            scheme1('valid')[i],
        ])

    for i in range(len(scheme2('rhoa'))):
        a = scheme2_dic[scheme2('a')[i]] + 1
        b = scheme2_dic[scheme2('b')[i]] + 1
        m = scheme2_dic[scheme2('m')[i]] + 1
        n = scheme2_dic[scheme2('n')[i]] + 1
        if [a, b, m, n] not in scheme_sensors:
            scheme_sensors.append([a, b, m, n])
            scheme_data.append([
                scheme2('err')[i],
                scheme2('i')[i],
                scheme2('ip')[i],
                scheme2('iperr')[i],
                scheme2('k')[i],
                scheme2('r')[i],
                scheme2('rhoa')[i],
                scheme2('u')[i],
                scheme2('valid')[i],
            ])
    logging.info('Merging complete (%d entries)', len(scheme_data))

    # Write to temporary file
    tmp_file = tmp_dir + 'tmpMerge.shm'
    f = open(tmp_file, 'w')
    # - electrodes
    f.write('%d\n' % len(electrodes))
    f.write('# x y z\n')
    for i in range(len(electrodes)):
        f.write('%f %f %f\n' %
                (electrodes[i][0], electrodes[i][1], electrodes[i][2]))
    # - data
    f.write('%d\n' % len(scheme_data))
    f.write('# a b m n err i ip iperr k r rhoa u valid \n')
    for i in range(len(scheme_data)):
        f.write('%d %d %d %d %f %f %f %f %f %f %f %f %d\n' %
                (scheme_sensors[i][0], scheme_sensors[i][1],
                 scheme_sensors[i][2], scheme_sensors[i][3], scheme_data[i][0],
                 scheme_data[i][1], scheme_data[i][2], scheme_data[i][3],
                 scheme_data[i][4], scheme_data[i][5], scheme_data[i][6],
                 scheme_data[i][7], scheme_data[i][8]))
    f.write('0\n')
    f.close()
    # Read file again to not generate any inconsistencies
    result = pb.load(tmp_file)
    if remove_tmp_file:
        os.remove(tmp_file)
    return result