angus-g / pyrol

Python interface to ROL 2.0
GNU Lesser General Public License v3.0
2 stars 2 forks source link

Error : Unable to recognise "Line search" #3

Open acse-yw11823 opened 1 week ago

acse-yw11823 commented 1 week ago

Hello. I am currently using Mac M1. And running the following code. It shows the error from the pybind11 ROL.Algorithm("Line Search", params). Does anyone know why"Line search" be recognised as string instead of ROL::Step? Does anyone have a hint about solving it? Her is my code:

import sys
sys.path.append("/home/spyro")

import os
os.environ["OMP_NUM_THREADS"] = "1"

from firedrake import *
import numpy as np
import finat
from ROL.firedrake_vector import FiredrakeVector as FeVector
import ROL
from mpi4py import MPI
import psutil

import spyro

sou_pos = [(0, 0.5)]
rec_pos = spyro.create_transect((-1.0, 0.25), (-1.0, 0.75), 10)

model = {}
model["opts"] = {
    "method": "KMV", 
    "quadrature": "KMV", 
    "degree": 1, 
    "dimension": 2, 
}
model["parallelism"] = {
    "type": "spatial", 
}
model["mesh"] = {
    "Lz": 1.0, 
    "Lx": 1.0, 
    "Ly": 0.0, 
    "meshfile": "not_used.msh", 
    "initmodel": "not_used.hdf5", 
    "truemodel": "not_used.hdf5", 
}
model["BCs"] = {
    "status": True, 
    "outer_bc": "non-reflective", 
    "damping_type": "polynomial", 
    "exponent": 2, 
    "cmax": 4.5, 
    "R": 1e-6, 
    "lz": 0.25, 
    "lx": 0.25,  
    "ly": 0.0, 
}
model["acquisition"] = {
    "source_type": "Ricker", 
    "source_pos": sou_pos, 
    "frequency": 5.0, 
    "delay": 0.1, 
    "receiver_locations": rec_pos, 
}
model["timeaxis"] = {
    "t0": 0.0, 
    "tf": 2.0, 
    "dt": 5e-4, 
    "amplitude": 1, 
    "nspool": 100, 
    "fspool": 100, 
}

mesh = RectangleMesh(150, 150, 1.5, 1.5)
mesh.coordinates.dat.data[:, 0] -= 1.25
mesh.coordinates.dat.data[:, 1] -= 0.25

comm = spyro.utils.mpi_init(model)
element = spyro.domains.space.FE_method(mesh, model["opts"]["method"], model["opts"]["degree"])
V = FunctionSpace(mesh, element)

x, y = SpatialCoordinate(mesh)
velocity = conditional(x > -0.5, 1, 2)
vp = Function(V, name="velocity").interpolate(velocity)
if comm.ensemble_comm.rank == 0:
    File("simple_velocity_model.pvd", comm=comm.comm).write(vp)

sources = spyro.Sources(model, mesh, V, comm)
receivers = spyro.Receivers(model, mesh, V, comm)
wavelet = spyro.full_ricker_wavelet(
    dt=model["timeaxis"]["dt"],
    tf=model["timeaxis"]["tf"],
    freq=model["acquisition"]["frequency"],
)

p_field, p_at_recv = spyro.solvers.forward(model, mesh, comm, vp, sources, wavelet, receivers)
spyro.plots.plot_shots(model, comm, p_at_recv)
spyro.io.save_shots(model, comm, p_at_recv)

outdir = "out/"
if not os.path.exists(outdir):
    os.mkdir(outdir)

velocity = conditional(x < 0, 1.5, 1.5)
vp = Function(V, name="velocity").interpolate(velocity)
if comm.ensemble_comm.rank == 0:
    File("simple_velocity_guess.pvd", comm=comm.comm).write(vp)

if comm.ensemble_comm.rank == 0:
    control_file = File(outdir + "control.pvd", comm=comm.comm)
    grad_file = File(outdir + "grad.pvd", comm=comm.comm)
quad_rule = finat.quadrature.make_quadrature(V.finat_element.cell, V.ufl_element().degree(), "KMV")
dxlump = dx(scheme=quad_rule)

class L2Inner(object):
    def __init__(self):
        self.A = assemble(TrialFunction(V) * TestFunction(V) * dxlump, mat_type="matfree")
        self.Ap = as_backend_type(self.A).mat()

    def eval(self, _u, _v):
        upet = as_backend_type(_u).vec()
        vpet = as_backend_type(_v).vec()
        A_u = self.Ap.createVecLeft()
        self.Ap.mult(upet, A_u)
        return vpet.dot(A_u)

def regularize_gradient(vp, dJ):
    m_u = TrialFunction(V)
    m_v = TestFunction(V)
    mgrad = m_u * m_v * dx(scheme=quad_rule)
    ffG = dot(grad(vp), grad(m_v)) * dx(scheme=quad_rule)
    G = mgrad - ffG
    lhsG, rhsG = lhs(G), rhs(G)
    gradreg = Function(V)
    grad_prob = LinearVariationalProblem(lhsG, rhsG, gradreg)
    grad_solver = LinearVariationalSolver(
        grad_prob,
        solver_parameters={
            "ksp_type": "preonly",
            "pc_type": "jacobi",
            "mat_type": "matfree",
        },
    )
    grad_solver.solve()
    dJ += gradreg
    return dJ

class Objective(ROL.Objective):
    def __init__(self, inner_product):
        ROL.Objective.__init__(self)
        self.inner_product = inner_product
        self.p_guess = None
        self.misfit = 0.0
        self.p_exact_recv = spyro.io.load_shots(model, comm)

    def value(self, x, tol):
        J_total = np.zeros((1))
        self.p_guess, p_guess_recv = spyro.solvers.forward(
            model,
            mesh,
            comm,
            vp,
            sources,
            wavelet,
            receivers,
        )
        self.misfit = spyro.utils.evaluate_misfit(model, p_guess_recv, self.p_exact_recv)
        J_total[0] += spyro.utils.compute_functional(model, self.misfit, velocity=vp)
        J_total = COMM_WORLD.allreduce(J_total, op=MPI.SUM)
        J_total[0] /= comm.ensemble_comm.size
        if comm.comm.size > 1:
            J_total[0] /= comm.comm.size
        return J_total[0]

    def gradient(self, g, x, tol):
        dJ = Function(V, name="gradient")
        dJ_local = spyro.solvers.gradient(
            model,
            mesh,
            comm,
            vp,
            receivers,
            self.p_guess,
            self.misfit,
        )
        if comm.ensemble_comm.size > 1:
            comm.allreduce(dJ_local, dJ)
        else:
            dJ = dJ_local
        dJ /= comm.ensemble_comm.size
        if comm.comm.size > 1:
            dJ /= comm.comm.size
        if "regularization" in model["opts"] and model["opts"]["regularization"]:
            dJ = regularize_gradient(vp, dJ)
        if comm.ensemble_comm.rank == 0:
            grad_file.write(dJ)
        g.scale(0)
        g.vec += dJ

    def update(self, x, flag, iteration):
        vp.assign(Function(V, x.vec, name="velocity"))
        if iteration >= 0:
            if comm.ensemble_comm.rank == 0:
                control_file.write(vp)

paramsDict = {
    "General": {"Secant": {"Type": "Limited-Memory BFGS", "Maximum Storage": 10}},
    "Step": {
        "Type": "Augmented Lagrangian",
        "Augmented Lagrangian": {
            "Subproblem Step Type": "Line Search",
            "Subproblem Iteration Limit": 5,
        },
        "Line Search": {"Descent Method": {"Type": "Quasi-Newton Step"}},
    },
    "Status Test": {
        "Gradient Tolerance": 1e-15,
        "Iteration Limit": 100,
        "Step Tolerance": 1e-15,
    },
}

params = ROL.ParameterList(paramsDict, "Parameters")

inner_product = L2Inner()

obj = Objective(inner_product)

u = Function(V, name="velocity").assign(vp)

opt = FeVector(u.vector(), inner_product)

algo = ROL.Algorithm("Line Search", params)

algo.run(opt, obj)

if comm.ensemble_comm.rank == 0:
    File("res.pvd", comm=comm.comm).write(vp)

And this is the error:

---------------------------------------------------------------------------
TypeError                                 Traceback (most recent call last)
Cell In[2], [line 224](vscode-notebook-cell:?execution_count=2&line=224)
    [220](vscode-notebook-cell:?execution_count=2&line=220) u = Function(V, name="velocity").assign(vp)
    [222](vscode-notebook-cell:?execution_count=2&line=222) opt = FeVector(u.vector(), inner_product)
--> [224](vscode-notebook-cell:?execution_count=2&line=224) algo = ROL.Algorithm("Line Search", params)
    [226](vscode-notebook-cell:?execution_count=2&line=226) algo.run(opt, obj)
    [228](vscode-notebook-cell:?execution_count=2&line=228) if comm.ensemble_comm.rank == 0:

TypeError: __init__(): incompatible constructor arguments. The following argument types are supported:
    1. _ROL.Algorithm(arg0: ROL::Step<double>, arg1: _ROL.StatusTest)

Invoked with: 'Line Search', <_ROL.ParameterList object at 0x16c2aecf0>

Many thanks

sghelichkhani commented 6 days ago

There are so many layers involved in the code that is not really practical as a reproducer. Below should be enough as a reproducer:

import ROL
from firedrake import *

paramsDict = {
    "General": {"Secant": {"Type": "Limited-Memory BFGS", "Maximum Storage": 10}},
    "Step": {
        "Type": "Augmented Lagrangian",
        "Augmented Lagrangian": {
            "Subproblem Step Type": "Line Search",
            "Subproblem Iteration Limit": 5,
        },
        "Line Search": {"Descent Method": {"Type": "Quasi-Newton Step"}},
    },
    "Status Test": {
        "Gradient Tolerance": 1e-15,
        "Iteration Limit": 100,
        "Step Tolerance": 1e-15,
    },
}

params = ROL.ParameterList(paramsDict, "Parameters")
algo = ROL.Algorithm("Line Search", params)

So the two inputs to ROL.Algorithm are arg0: ROL::Step<double>, arg1: _ROL.StatusTest. So both of the arguments you are passing in are not the correct type.

The fact is, although ROL.Algorithm is an interface for the general roltrilinos algorithms class, but is only implemented by @angus-g to provide access to a few optimisation methods that have not been previously accessible within ROL2.0, like Lin-More, which seems to be the one optimisation method that rol's developers have focused on for their large-scale problems in recent years.

My suggestion is that either you follow the few demos that are available in pyadjoint/firedrake repositories that are generally based on the old ROL routines (pyROL2.0 is backward compatible) or if you want to take advantage of the new implementations, look at our test cases or demos in gadopt: https://github.com/g-adopt/g-adopt/blob/master/tests/optimisation_checkpointing/helmholtz.py