<reponame>ChuningWang/pyroms
import numpy as np
from scipy.optimize import linprog
from scipy.sparse import coo_matrix
from xarray import DataArray
from bathy_smoother import LP_bathy_tools, bathy_tools

'''
This code is adapted from the matlab code "LP Bathymetry" by <NAME> (http://drobilica.irb.hr/~mathieu/Bathymetry/index.html).
For a description of the method, see

<NAME>, <NAME>, <NAME>, A new approach to bathymetry
smoothing in sigma-coordinate ocean models, Ocean Modelling 29 (2009)
128--136.

*** Some modifications done by <NAME>

The MATLAB LP Bathymetry package is a very old package dated back to
2021. The computation involves a lot of loops and is not vectornized for
MATLAB. This makes it extremely slow for large grids.

The python code inherited all its disadvantages, and when I tested the
code it takes more than one day to process the IJS values when smoothing
a 1000x1000 grid. This is unbearable. So I decided to overhaul the
smoother to make it more Pythonish.

The overhaul consists of two parts.

Part 1, rewrite the tool functions such as GetIJS_**. Rewrite the loops
in the code using numpy array operations. The Neighborhood finder
function is also rewritten, but the 'ConnectedComponent' function is
untouched.

Part 2, update the linear programming solver. The original linear
programming solver, lpsolve55 keeps crashing when I run the code. Also,
compile and installing the package can be a problem for new Python users.
In order to simplify it (and also take advantage of recent development)
in python scientific computation, I replaced lpsolve55 with the scipy
linear programming solver (scipy.optimize.linprog). linprog is pretty
fast and support sparse matrix, which is memory conserving.

Since I cannot run lpsolve on my machine I don't know if the two solvers
produce the exact same result. But the smoothed bathymetry looks
reasonable, and it also satisfies the rx0 criterion, which I consider a
success.

2021/9/14
'''


def LP_smoothing_rx0(MSK, Hobs, rx0max, SignConst=0, AmpConst=1.e5,
                     verbose=True):
    """
    This program perform a linear programming method in order to
    optimize the bathymetry for a fixed factor r.
    The inequality |H(e)-H(e')| / (H(e)-H(e')) <= r where H(e)=h(e)+dh(e)
    can be rewritten as two linear inequalities on dh(e) and dh(e').
    The optimal bathymetry is obtain by minimising the perturbation
    P = sum_e(|dh(e)| under the above inequalitie constraintes.

    Usage:
    Hnew = LP_smoothing_rx0(MSK, Hobs, rx0max, SignConst, AmpConst)

    ---MSK(eta_rho,xi_rho) is the mask of the grd
         1 for sea
         0 for land
    ---Hobs(eta_rho,xi_rho) is the raw depth of the grid
    ---rx0max is the target rx0 roughness factor
    ---SignConst, single value or (eta_rho,xi_rho) matrix of 0, +1, -1
         +1  only bathymetry increase are allowed.
         -1  only bathymetry decrease are allowed.
         0   increase and decrease are allowed.
         (put 0 if you are indifferent)
    ---AmpConst, single value or (eta_rho,xi_rho)  matrix of reals.
         coefficient alpha such that the new bathymetry should
         satisfy to  |h^{new} - h^{raw}| <= alpha h^{raw}
         (put 10000 if you are indifferent)
    """

    use_xarray = False
    if isinstance(Hobs, DataArray):
        use_xarray = True
        coords = Hobs.coords
        Hobs = Hobs.values
    if isinstance(MSK, DataArray):
        MSK = MSK.values

    # Get the I, J indices and coefficients for the solver
    eta_rho, xi_rho = MSK.shape

    iList, jList, sList, Constant = \
        LP_bathy_tools.GetIJS_rx0(MSK, Hobs, rx0max, verbose=verbose)

    iListApp, jListApp, sListApp, ConstantApp = \
        LP_bathy_tools.GetIJS_maxamp(MSK, Hobs, AmpConst, verbose=verbose)

    iList, jList, sList, Constant = \
        LP_bathy_tools.MergeIJS_listings(iList, jList, sList, Constant,
                                         iListApp, jListApp, sListApp,
                                         ConstantApp)

    iListApp, jListApp, sListApp, ConstantApp = \
        LP_bathy_tools.GetIJS_signs(MSK, SignConst, verbose=verbose)

    iList, jList, sList, Constant = \
        LP_bathy_tools.MergeIJS_listings(iList, jList, sList, Constant,
                                         iListApp, jListApp, sListApp,
                                         ConstantApp)

    TotalNbVert = int(MSK.sum())

    ObjectiveFct = np.zeros((2*TotalNbVert))
    for iVert in range(TotalNbVert):
        ObjectiveFct[TotalNbVert+iVert] = 1

    iList = iList.astype(int)
    jList = jList.astype(int)

    # Build the sparse matrix and solver
    shp = (int(len(jList)/2), 2*TotalNbVert)
    mtx = coo_matrix((sList, (iList-1, jList-1)), shape=shp)
    print('Solver Matrix Size, %d5, %d5' % shp)
    solver = linprog(c=ObjectiveFct, A_ub=mtx, b_ub=Constant,
                     options={'sparse': True})

    if not solver['success']:
        raise ValueError('Linear programming solver failed.')
    hcorr = solver['x']

    Hnew = Hobs.copy()
    Hnew[MSK == 1] = Hobs[MSK == 1] + hcorr[:TotalNbVert]
    MaxRx0 = bathy_tools.RoughnessMatrix(Hnew, MSK).max()
    if verbose:
        print('rx0max = ', rx0max, '  MaxRx0 = ', MaxRx0)

    if use_xarray:
        Hnew = DataArray(Hnew, coords)

    return Hnew


def LP_smoothing_rx0_heuristic(MSK, Hobs, rx0max, SignConst=0, AmpConst=1.e5):
    """
    This program perform a linear programming method in order to
    optimize the bathymetry for a fixed factor r.
    The inequality |H(e)-H(e')| / (H(e)-H(e')) <= r where H(e)=h(e)+dh(e)
    can be rewritten as two linear inequalities on dh(e) and dh(e').
    The optimal bathymetry is obtain by minimising the perturbation
    P = sum_e(|dh(e)| under the above inequalitie constraintes.
    In order to reduce the computation time, an heuristic method is
    used.

    Usage:
    Hnew = LP_smoothing_rx0_heuristic(MSK, Hobs, rx0max, SignConst,
                                          AmpConst)

    ---MSK(eta_rho,xi_rho) is the mask of the grd
         1 for sea
         0 for land
    ---Hobs(eta_rho,xi_rho) is the raw depth of the grid
    ---rx0max is the target rx0 roughness factor
    ---SignConst, single value or (eta_rho,xi_rho) matrix of 0, +1, -1
         +1  only bathymetry increase are allowed.
         -1  only bathymetry decrease are allowed.
         0   increase and decrease are allowed.
         (put 0 if you are indifferent)
    ---AmpConst, single value or (eta_rho,xi_rho)  matrix of reals.
         coefficient alpha such that the new bathymetry should
         satisfy to  |h^{new} - h^{raw}| <= alpha h^{raw}
         (put 10000 if you are indifferent)
    """

    use_xarray = False
    if isinstance(Hobs, DataArray):
        use_xarray = True
        coords = Hobs.coords
        Hobs = Hobs.values
    if isinstance(MSK, DataArray):
        MSK = MSK.values

    # the points that need to be modified
    MSKbad = LP_bathy_tools.GetBadPoints(MSK, Hobs, rx0max)

    eta_rho, xi_rho = MSK.shape

    Kdist = 5

    Kbad = np.where(MSKbad == 1)
    nbKbad = np.size(Kbad, 1)
    ListIdx = np.zeros((eta_rho, xi_rho), dtype=np.int)
    ListIdx[Kbad] = list(range(nbKbad))

    ListEdges = []
    nbEdge = 0
    for iK in range(nbKbad):
        iEta, iXi = Kbad[0][iK], Kbad[1][iK]
        ListNeigh = LP_bathy_tools.Neighborhood(MSK, iEta, iXi, 2*Kdist+1)
        bidx = np.where(MSKbad[ListNeigh[:, 0], ListNeigh[:, 1]] == 1)
        idx = ListIdx[ListNeigh[bidx, 0], ListNeigh[bidx, 1]]
        idx = idx[idx > iK]
        for iidx in idx:
            ListEdges.append([iK, iidx])
        nbEdge = nbEdge + len(idx)

    ListEdges = np.array(ListEdges)
    ListVertexStatus = LP_bathy_tools.ConnectedComponent(ListEdges, nbKbad)
    nbColor = int(ListVertexStatus.max())

    Hnew = Hobs.copy()
    print('-------------------------------------------------------')
    print('Total iColor, %10d', nbColor)
    for iColor in range(1, nbColor+1):
        MSKcolor = np.zeros((eta_rho, xi_rho))
        K = np.where(ListVertexStatus == iColor)
        nbK = np.size(K, 1)
        print('iColor = %10d, nbK = %6d' % (iColor, nbK), end='\r')
        for iVertex in range(nbKbad):
            if (ListVertexStatus[iVertex] == iColor):
                iEta, iXi = Kbad[0][iVertex], Kbad[1][iVertex]
                MSKcolor[iEta, iXi] = 1
                ListNeigh = LP_bathy_tools.Neighborhood(MSK, iEta, iXi, Kdist)
                nbNeigh = np.size(ListNeigh, 0)
                for iNeigh in range(nbNeigh):
                    iEtaN, iXiN = ListNeigh[iNeigh]
                    MSKcolor[iEtaN, iXiN] = 1
        K = np.where(MSKcolor == 1)
        MSKHobs = np.zeros((eta_rho, xi_rho))
        MSKHobs[K] = Hobs[K].copy()
        HnewI = \
            LP_smoothing_rx0(MSKcolor, MSKHobs, rx0max, SignConst, AmpConst,
                             verbose=False)
        Hnew[K] = HnewI[K]

    print('Final obtained bathymetry')
    RMat = bathy_tools.RoughnessMatrix(Hnew, MSK)
    MaxRx0 = RMat.max()
    print('rx0max = ', rx0max, '  MaxRx0 = ', MaxRx0)

    if use_xarray:
        Hnew = DataArray(Hnew, coords)

    return Hnew

import numpy as np
from scipy.optimize import linprog
from scipy.sparse import coo_matrix
from xarray import DataArray
from bathy_smoother import LP_bathy_tools, bathy_tools

'''
This code is adapted from the matlab code "LP Bathymetry" by <NAME> (http://drobilica.irb.hr/~mathieu/Bathymetry/index.html).
For a description of the method, see

<NAME>, <NAME>, <NAME>, A new approach to bathymetry
smoothing in sigma-coordinate ocean models, Ocean Modelling 29 (2009)
128--136.

*** Some modifications done by <NAME>

The MATLAB LP Bathymetry package is a very old package dated back to
2021. The computation involves a lot of loops and is not vectornized for
MATLAB. This makes it extremely slow for large grids.

The python code inherited all its disadvantages, and when I tested the
code it takes more than one day to process the IJS values when smoothing
a 1000x1000 grid. This is unbearable. So I decided to overhaul the
smoother to make it more Pythonish.

The overhaul consists of two parts.

Part 1, rewrite the tool functions such as GetIJS_**. Rewrite the loops
in the code using numpy array operations. The Neighborhood finder
function is also rewritten, but the 'ConnectedComponent' function is
untouched.

Part 2, update the linear programming solver. The original linear
programming solver, lpsolve55 keeps crashing when I run the code. Also,
compile and installing the package can be a problem for new Python users.
In order to simplify it (and also take advantage of recent development)
in python scientific computation, I replaced lpsolve55 with the scipy
linear programming solver (scipy.optimize.linprog). linprog is pretty
fast and support sparse matrix, which is memory conserving.

Since I cannot run lpsolve on my machine I don't know if the two solvers
produce the exact same result. But the smoothed bathymetry looks
reasonable, and it also satisfies the rx0 criterion, which I consider a
success.

2021/9/14
'''


def LP_smoothing_rx0(MSK, Hobs, rx0max, SignConst=0, AmpConst=1.e5,
                     verbose=True):
    """
    This program perform a linear programming method in order to
    optimize the bathymetry for a fixed factor r.
    The inequality |H(e)-H(e')| / (H(e)-H(e')) <= r where H(e)=h(e)+dh(e)
    can be rewritten as two linear inequalities on dh(e) and dh(e').
    The optimal bathymetry is obtain by minimising the perturbation
    P = sum_e(|dh(e)| under the above inequalitie constraintes.

    Usage:
    Hnew = LP_smoothing_rx0(MSK, Hobs, rx0max, SignConst, AmpConst)

    ---MSK(eta_rho,xi_rho) is the mask of the grd
         1 for sea
         0 for land
    ---Hobs(eta_rho,xi_rho) is the raw depth of the grid
    ---rx0max is the target rx0 roughness factor
    ---SignConst, single value or (eta_rho,xi_rho) matrix of 0, +1, -1
         +1  only bathymetry increase are allowed.
         -1  only bathymetry decrease are allowed.
         0   increase and decrease are allowed.
         (put 0 if you are indifferent)
    ---AmpConst, single value or (eta_rho,xi_rho)  matrix of reals.
         coefficient alpha such that the new bathymetry should
         satisfy to  |h^{new} - h^{raw}| <= alpha h^{raw}
         (put 10000 if you are indifferent)
    """

    use_xarray = False
    if isinstance(Hobs, DataArray):
        use_xarray = True
        coords = Hobs.coords
        Hobs = Hobs.values
    if isinstance(MSK, DataArray):
        MSK = MSK.values

    # Get the I, J indices and coefficients for the solver
    eta_rho, xi_rho = MSK.shape

    iList, jList, sList, Constant = \
        LP_bathy_tools.GetIJS_rx0(MSK, Hobs, rx0max, verbose=verbose)

    iListApp, jListApp, sListApp, ConstantApp = \
        LP_bathy_tools.GetIJS_maxamp(MSK, Hobs, AmpConst, verbose=verbose)

    iList, jList, sList, Constant = \
        LP_bathy_tools.MergeIJS_listings(iList, jList, sList, Constant,
                                         iListApp, jListApp, sListApp,
                                         ConstantApp)

    iListApp, jListApp, sListApp, ConstantApp = \
        LP_bathy_tools.GetIJS_signs(MSK, SignConst, verbose=verbose)

    iList, jList, sList, Constant = \
        LP_bathy_tools.MergeIJS_listings(iList, jList, sList, Constant,
                                         iListApp, jListApp, sListApp,
                                         ConstantApp)

    TotalNbVert = int(MSK.sum())

    ObjectiveFct = np.zeros((2*TotalNbVert))
    for iVert in range(TotalNbVert):
        ObjectiveFct[TotalNbVert+iVert] = 1

    iList = iList.astype(int)
    jList = jList.astype(int)

    # Build the sparse matrix and solver
    shp = (int(len(jList)/2), 2*TotalNbVert)
    mtx = coo_matrix((sList, (iList-1, jList-1)), shape=shp)
    print('Solver Matrix Size, %d5, %d5' % shp)
    solver = linprog(c=ObjectiveFct, A_ub=mtx, b_ub=Constant,
                     options={'sparse': True})

    if not solver['success']:
        raise ValueError('Linear programming solver failed.')
    hcorr = solver['x']

    Hnew = Hobs.copy()
    Hnew[MSK == 1] = Hobs[MSK == 1] + hcorr[:TotalNbVert]
    MaxRx0 = bathy_tools.RoughnessMatrix(Hnew, MSK).max()
    if verbose:
        print('rx0max = ', rx0max, '  MaxRx0 = ', MaxRx0)

    if use_xarray:
        Hnew = DataArray(Hnew, coords)

    return Hnew


def LP_smoothing_rx0_heuristic(MSK, Hobs, rx0max, SignConst=0, AmpConst=1.e5):
    """
    This program perform a linear programming method in order to
    optimize the bathymetry for a fixed factor r.
    The inequality |H(e)-H(e')| / (H(e)-H(e')) <= r where H(e)=h(e)+dh(e)
    can be rewritten as two linear inequalities on dh(e) and dh(e').
    The optimal bathymetry is obtain by minimising the perturbation
    P = sum_e(|dh(e)| under the above inequalitie constraintes.
    In order to reduce the computation time, an heuristic method is
    used.

    Usage:
    Hnew = LP_smoothing_rx0_heuristic(MSK, Hobs, rx0max, SignConst,
                                          AmpConst)

    ---MSK(eta_rho,xi_rho) is the mask of the grd
         1 for sea
         0 for land
    ---Hobs(eta_rho,xi_rho) is the raw depth of the grid
    ---rx0max is the target rx0 roughness factor
    ---SignConst, single value or (eta_rho,xi_rho) matrix of 0, +1, -1
         +1  only bathymetry increase are allowed.
         -1  only bathymetry decrease are allowed.
         0   increase and decrease are allowed.
         (put 0 if you are indifferent)
    ---AmpConst, single value or (eta_rho,xi_rho)  matrix of reals.
         coefficient alpha such that the new bathymetry should
         satisfy to  |h^{new} - h^{raw}| <= alpha h^{raw}
         (put 10000 if you are indifferent)
    """

    use_xarray = False
    if isinstance(Hobs, DataArray):
        use_xarray = True
        coords = Hobs.coords
        Hobs = Hobs.values
    if isinstance(MSK, DataArray):
        MSK = MSK.values

    # the points that need to be modified
    MSKbad = LP_bathy_tools.GetBadPoints(MSK, Hobs, rx0max)

    eta_rho, xi_rho = MSK.shape

    Kdist = 5

    Kbad = np.where(MSKbad == 1)
    nbKbad = np.size(Kbad, 1)
    ListIdx = np.zeros((eta_rho, xi_rho), dtype=np.int)
    ListIdx[Kbad] = list(range(nbKbad))

    ListEdges = []
    nbEdge = 0
    for iK in range(nbKbad):
        iEta, iXi = Kbad[0][iK], Kbad[1][iK]
        ListNeigh = LP_bathy_tools.Neighborhood(MSK, iEta, iXi, 2*Kdist+1)
        bidx = np.where(MSKbad[ListNeigh[:, 0], ListNeigh[:, 1]] == 1)
        idx = ListIdx[ListNeigh[bidx, 0], ListNeigh[bidx, 1]]
        idx = idx[idx > iK]
        for iidx in idx:
            ListEdges.append([iK, iidx])
        nbEdge = nbEdge + len(idx)

    ListEdges = np.array(ListEdges)
    ListVertexStatus = LP_bathy_tools.ConnectedComponent(ListEdges, nbKbad)
    nbColor = int(ListVertexStatus.max())

    Hnew = Hobs.copy()
    print('-------------------------------------------------------')
    print('Total iColor, %10d', nbColor)
    for iColor in range(1, nbColor+1):
        MSKcolor = np.zeros((eta_rho, xi_rho))
        K = np.where(ListVertexStatus == iColor)
        nbK = np.size(K, 1)
        print('iColor = %10d, nbK = %6d' % (iColor, nbK), end='\r')
        for iVertex in range(nbKbad):
            if (ListVertexStatus[iVertex] == iColor):
                iEta, iXi = Kbad[0][iVertex], Kbad[1][iVertex]
                MSKcolor[iEta, iXi] = 1
                ListNeigh = LP_bathy_tools.Neighborhood(MSK, iEta, iXi, Kdist)
                nbNeigh = np.size(ListNeigh, 0)
                for iNeigh in range(nbNeigh):
                    iEtaN, iXiN = ListNeigh[iNeigh]
                    MSKcolor[iEtaN, iXiN] = 1
        K = np.where(MSKcolor == 1)
        MSKHobs = np.zeros((eta_rho, xi_rho))
        MSKHobs[K] = Hobs[K].copy()
        HnewI = \
            LP_smoothing_rx0(MSKcolor, MSKHobs, rx0max, SignConst, AmpConst,
                             verbose=False)
        Hnew[K] = HnewI[K]

    print('Final obtained bathymetry')
    RMat = bathy_tools.RoughnessMatrix(Hnew, MSK)
    MaxRx0 = RMat.max()
    print('rx0max = ', rx0max, '  MaxRx0 = ', MaxRx0)

    if use_xarray:
        Hnew = DataArray(Hnew, coords)

    return Hnew
This code is adapted from the matlab code "LP Bathymetry" by <NAME> (http://drobilica.irb.hr/~mathieu/Bathymetry/index.html). For a description of the method, see  <NAME>, <NAME>, <NAME>, A new approach to bathymetry smoothing in sigma-coordinate ocean models, Ocean Modelling 29 (2009) 128--136.  *** Some modifications done by <NAME>  The MATLAB LP Bathymetry package is a very old package dated back to 2021. The computation involves a lot of loops and is not vectornized for MATLAB. This makes it extremely slow for large grids.  The python code inherited all its disadvantages, and when I tested the code it takes more than one day to process the IJS values when smoothing a 1000x1000 grid. This is unbearable. So I decided to overhaul the smoother to make it more Pythonish.  The overhaul consists of two parts.  Part 1, rewrite the tool functions such as GetIJS_**. Rewrite the loops in the code using numpy array operations. The Neighborhood finder function is also rewritten, but the 'ConnectedComponent' function is untouched.  Part 2, update the linear programming solver. The original linear programming solver, lpsolve55 keeps crashing when I run the code. Also, compile and installing the package can be a problem for new Python users. In order to simplify it (and also take advantage of recent development) in python scientific computation, I replaced lpsolve55 with the scipy linear programming solver (scipy.optimize.linprog). linprog is pretty fast and support sparse matrix, which is memory conserving.  Since I cannot run lpsolve on my machine I don't know if the two solvers produce the exact same result. But the smoothed bathymetry looks reasonable, and it also satisfies the rx0 criterion, which I consider a success.  2021/9/14 This program perform a linear programming method in order to     optimize the bathymetry for a fixed factor r.     The inequality |H(e)-H(e')| / (H(e)-H(e')) <= r where H(e)=h(e)+dh(e)     can be rewritten as two linear inequalities on dh(e) and dh(e').     The optimal bathymetry is obtain by minimising the perturbation     P = sum_e(|dh(e)| under the above inequalitie constraintes.      Usage:     Hnew = LP_smoothing_rx0(MSK, Hobs, rx0max, SignConst, AmpConst)      ---MSK(eta_rho,xi_rho) is the mask of the grd          1 for sea          0 for land     ---Hobs(eta_rho,xi_rho) is the raw depth of the grid     ---rx0max is the target rx0 roughness factor     ---SignConst, single value or (eta_rho,xi_rho) matrix of 0, +1, -1          +1  only bathymetry increase are allowed.          -1  only bathymetry decrease are allowed.          0   increase and decrease are allowed.          (put 0 if you are indifferent)     ---AmpConst, single value or (eta_rho,xi_rho)  matrix of reals.          coefficient alpha such that the new bathymetry should          satisfy to  |h^{new} - h^{raw}| <= alpha h^{raw}          (put 10000 if you are indifferent) # Get the I, J indices and coefficients for the solver # Build the sparse matrix and solver This program perform a linear programming method in order to     optimize the bathymetry for a fixed factor r.     The inequality |H(e)-H(e')| / (H(e)-H(e')) <= r where H(e)=h(e)+dh(e)     can be rewritten as two linear inequalities on dh(e) and dh(e').     The optimal bathymetry is obtain by minimising the perturbation     P = sum_e(|dh(e)| under the above inequalitie constraintes.     In order to reduce the computation time, an heuristic method is     used.      Usage:     Hnew = LP_smoothing_rx0_heuristic(MSK, Hobs, rx0max, SignConst,                                           AmpConst)      ---MSK(eta_rho,xi_rho) is the mask of the grd          1 for sea          0 for land     ---Hobs(eta_rho,xi_rho) is the raw depth of the grid     ---rx0max is the target rx0 roughness factor     ---SignConst, single value or (eta_rho,xi_rho) matrix of 0, +1, -1          +1  only bathymetry increase are allowed.          -1  only bathymetry decrease are allowed.          0   increase and decrease are allowed.          (put 0 if you are indifferent)     ---AmpConst, single value or (eta_rho,xi_rho)  matrix of reals.          coefficient alpha such that the new bathymetry should          satisfy to  |h^{new} - h^{raw}| <= alpha h^{raw}          (put 10000 if you are indifferent) # the points that need to be modified