How to use the devito.Function function in devito

coordinates=wave.geometry.rec_positions)

        gradient, _ = wave.jacobian_adjoint(residual, u0, vp=v0,
                                            checkpointing=checkpointing)
        G = np.dot(gradient.data.reshape(-1), dm.reshape(-1))

        # FWI Gradient test
        H = [0.5, 0.25, .125, 0.0625, 0.0312, 0.015625, 0.0078125]
        error1 = np.zeros(7)
        error2 = np.zeros(7)
        for i in range(0, 7):
            # Add the perturbation to the model
            def initializer(data):
                data[:] = np.sqrt(v0.data**2 * v**2 /
                                  ((1 - H[i]) * v**2 + H[i] * v0.data**2))
            vloc = Function(name='vloc', grid=wave.model.grid, space_order=space_order,
                            initializer=initializer)
            # Data for the new model
            d = wave.forward(vp=vloc)[0]
            # First order error Phi(m0+dm) - Phi(m0)
            F_i = .5*linalg.norm((d.data - rec.data).reshape(-1))**2
            error1[i] = np.absolute(F_i - F0)
            # Second order term r Phi(m0+dm) - Phi(m0) - 
            error2[i] = np.absolute(F_i - F0 - H[i] * G)

        # Test slope of the  tests
        p1 = np.polyfit(np.log10(H), np.log10(error1), 1)
        p2 = np.polyfit(np.log10(H), np.log10(error2), 1)
        info('1st order error, Phi(m0+dm)-Phi(m0): %s' % (p1))
        info(r'2nd order error, Phi(m0+dm)-Phi(m0) - : %s' % (p2))
        assert np.isclose(p1[0], 1.0, rtol=0.1)
        assert np.isclose(p2[0], 2.0, rtol=0.1)

    @pytest.mark.parametrize('eq,expected,blocking', [
        ('Eq(f, 2*f)', [2, 0, 0], False),
        ('Eq(u, 2*u)', [0, 2, 0, 0], False),
        ('Eq(u, 2*u)', [3, 0, 0, 0, 0, 0], True)
    ])
    @patch("devito.passes.iet.openmp.Ompizer.COLLAPSE_NCORES", 1)
    @patch("devito.passes.iet.openmp.Ompizer.COLLAPSE_WORK", 0)
    def test_collapsing(self, eq, expected, blocking):
        grid = Grid(shape=(3, 3, 3))

        f = Function(name='f', grid=grid)  # noqa
        u = TimeFunction(name='u', grid=grid)  # noqa

        eq = eval(eq)

        if blocking:
            op = Operator(eq, dle=('blocking', 'simd', 'openmp', {'blockinner': True}))
            iterations = FindNodes(Iteration).visit(op._func_table['bf0'])
        else:
            op = Operator(eq, dle=('simd', 'openmp'))
            iterations = FindNodes(Iteration).visit(op)

        assert len(iterations) == len(expected)

        # Check for presence of pragma omp + collapse clause
        for i, j in zip(iterations, expected):
            if j > 0:

    @pytest.mark.parallel(mode=4)
    def test_from_replicated_to_distributed(self):
        shape = (4, 4)
        grid = Grid(shape=shape)
        x, y = grid.dimensions
        glb_pos_map = grid.distributor.glb_pos_map
        u = Function(name='u', grid=grid, space_order=0)  # distributed
        v = Function(name='v', grid=grid, space_order=0)  # distributed
        a = np.arange(16).reshape(shape)  # replicated

        # Full array
        u.data[:] = a
        if LEFT in glb_pos_map[x] and LEFT in glb_pos_map[y]:
            assert np.all(u.data == [[0, 1], [4, 5]])
        elif LEFT in glb_pos_map[x] and RIGHT in glb_pos_map[y]:
            assert np.all(u.data == [[2, 3], [6, 7]])
        elif RIGHT in glb_pos_map[x] and LEFT in glb_pos_map[y]:
            assert np.all(u.data == [[8, 9], [12, 13]])
        else:
            assert np.all(u.data == [[10, 11], [14, 15]])

        # Subsection (all ranks touched)
        u.data[:] = 0
        u.data[1:3, 1:3] = a[1:3, 1:3]

def test_function():
    grid = Grid(shape=(3, 3, 3))
    f = Function(name='f', grid=grid)
    f.data[0] = 1.

    pkl_f = pickle.dumps(f)
    new_f = pickle.loads(pkl_f)

    # .data is initialized, so it should have been pickled too
    assert np.all(f.data[0] == 1.)
    assert np.all(new_f.data[0] == 1.)

    assert f.space_order == new_f.space_order
    assert f.dtype == new_f.dtype
    assert f.shape == new_f.shape

def test_haloupdate_multi_op(self):
        """
        Test that halo updates are carried out correctly when multiple operators
        are applied consecutively.
        """
        a = np.arange(64).reshape((8, 8))
        grid = Grid(shape=a.shape, extent=(8, 8))

        so = 3
        dims = grid.dimensions
        f = Function(name='f', grid=grid, space_order=so)
        f.data[:] = a

        fo = Function(name='fo', grid=grid, space_order=so)

        for d in dims:
            rhs = generic_derivative(f, d, so, 1)
            expr = Eq(fo, rhs)
            op = Operator(expr)
            op.apply()
            f.data[:, :] = fo.data[:, :]

        assert (np.isclose(norm(f), 17.24904, atol=1e-4, rtol=0))

return field.dx(x0=x-x.spacing/2)


def g2_tilde(field):
    return field.dy(x0=y-y.spacing/2)


def g3_tilde(field):
    return field.dz(x0=z-z.spacing/2)


# Functions  for additional factorization
b1mf = Function(name='b1mf', grid=grid, space_order=space_order)
b1m2e = Function(name='b1m2e', grid=grid, space_order=space_order)
b1mfa2 = Function(name='b1mfa2', grid=grid, space_order=space_order)
b1mfpfa2 = Function(name='b1mfpfa2', grid=grid, space_order=space_order)
bfes1ma2 = Function(name='bfes1ma2', grid=grid, space_order=space_order)

# Equations for additional factorization
eq_b1mf = Eq(b1mf, b * (1 - f))
eq_b1m2e = Eq(b1m2e, b * (1 + 2 * eps))
eq_b1mfa2 = Eq(b1mfa2, b * (1 - f * eta**2))
eq_b1mfpfa2 = Eq(b1mfpfa2, b * (1 - f + f * eta**2))
eq_bfes1ma2 = Eq(bfes1ma2, b * f * eta * sqrt(1 - eta**2))

# Time update equation for quasi-P state variable p
update_p_nl = t.spacing**2 * vel**2 / b * \
    (g1_tilde(b1m2e * g1(p0)) +
     g2_tilde(b1m2e * g2(p0)) +
     g3_tilde(b1mfa2 * g3(p0) + bfes1ma2 * g3(m0))) + \
    (2 - t.spacing * wOverQ) * p0 + (t.spacing * wOverQ - 1) * p0.backward

dtype = np.float32
npad = 20
qmin = 0.1
qmax = 1000.0
tmax = 20.0
fpeak = 0.010
omega = 2.0 * np.pi * fpeak

shape = (1201, 1201, 601)
spacing = (10.0, 10.0, 10.0)
origin = tuple([0.0 for s in shape])
extent = tuple([d * (s - 1) for s, d in zip(shape, spacing)])
grid = Grid(extent=extent, shape=shape, origin=origin, dtype=dtype)

b = Function(name='b', grid=grid, space_order=space_order)
f = Function(name='f', grid=grid, space_order=space_order)
vel = Function(name='vel', grid=grid, space_order=space_order)
eps = Function(name='eps', grid=grid, space_order=space_order)
eta = Function(name='eta', grid=grid, space_order=space_order)
wOverQ = Function(name='wOverQ', grid=grid, space_order=space_order)
theta = Function(name='theta', grid=grid, space_order=space_order)
phi = Function(name='phi', grid=grid, space_order=space_order)

b.data[:] = 1.0
f.data[:] = 0.84
vel.data[:] = 1.5
eps.data[:] = 0.2
eta.data[:] = 0.4
wOverQ.data[:] = 1.0
theta.data[:] = np.pi / 3
phi.data[:] = np.pi / 6

# Functions for additional factorization
b1p2e = Function(name='b1p2e', grid=grid, space_order=space_order)
b1mf = Function(name='b1mf', grid=grid, space_order=space_order)
b2epfa2 = Function(name='b2epfa2', grid=grid, space_order=space_order)
bfes1ma2 = Function(name='bfes1ma2', grid=grid, space_order=space_order)
bfa2 = Function(name='bfa2', grid=grid, space_order=space_order)

# Equations for additional factorization
eq_b1p2e = Eq(b1p2e, b * (1 + 2 * eps))
eq_b1mf = Eq(b1mf, b * (1 - f))
eq_b2epfa2 = Eq(b2epfa2, b * (2 * eps + f * eta**2))
eq_bfes1ma2 = Eq(bfes1ma2, b * f * eta * sqrt(1 - eta**2))
eq_bfa2 = Eq(bfa2, b * f * eta**2)

p_xyz = Function(name='p_xyz', grid=grid, space_order=space_order)
m_xyz = Function(name='m_xyz', grid=grid, space_order=space_order)

eq_pxyz = Eq(p_xyz, 
             my_d_tilde(b1p2e * my_d(p0, x), x) +
             my_d_tilde(b1p2e * my_d(p0, y), y) +
             my_d_tilde(b1p2e * my_d(p0, z), z) +
             g3_tilde(- b2epfa2 * g3(p0, phi, theta) +
                      bfes1ma2 * g3(m0, phi, theta), phi, theta), phi, theta)

eq_mxyz = Eq(m_xyz, 
             my_d_tilde(b1mf * my_d(m0, x), x) +
             my_d_tilde(b1mf * my_d(m0, y), y) +
             my_d_tilde(b1mf * my_d(m0, z), z) + 
             g3_tilde(+ bfes1ma2 * g3(p0, phi, theta) + 
                      bfa2 * g3(m0, phi, theta),  phi, theta), phi, theta)

def run(shape=(50, 50, 50), spacing=(10.0, 10.0, 10.0), tn=1000.0,
        space_order=4, nbl=40, full_run=False, autotune=False, **kwargs):

    solver = acoustic_sa_setup(shape=shape, spacing=spacing, nbl=nbl, tn=tn,
                               space_order=space_order, **kwargs)

    info("Applying Forward")
    # Define receiver geometry (spread across x, just below surface)
    rec, u, summary = solver.forward(save=full_run, autotune=autotune)

    if not full_run:
        return summary.gflopss, summary.oi, summary.timings, [rec, u.data]

    # Smooth velocity
    initial_vp = Function(name='v0', grid=solver.model.grid, space_order=space_order)
    smooth(initial_vp, solver.model.vp)
    dm = solver.model.vp - initial_vp

    info("Applying Adjoint")
    solver.adjoint(rec, autotune=autotune)
    info("Applying Born")
    solver.jacobian(dm, autotune=autotune)
    info("Applying Gradient")
    solver.jacobian_adjoint(rec, u, autotune=autotune)
    return summary.gflopss, summary.oi, summary.timings, [rec, u.data]

return field.dz(x0=z+z.spacing/2)


def g1_tilde(field):
    return field.dx(x0=x-x.spacing/2)


def g2_tilde(field):
    return field.dy(x0=y-y.spacing/2)


def g3_tilde(field):
    return field.dz(x0=z-z.spacing/2)


p_x = Function(name='p_x', grid=grid, space_order=space_order)
p_y = Function(name='p_y', grid=grid, space_order=space_order)
p_z = Function(name='p_z', grid=grid, space_order=space_order)

update_px = Eq(p_x, g1_tilde(b * g1(p0)))
update_py = Eq(p_y, g2_tilde(b * g2(p0)))
update_pz = Eq(p_z, g3_tilde(b * g3(p0)))

update_p0 = t.spacing**2 * vel**2 / b * \
    (p_x + p_y + p_z) + \
    (2 - t.spacing * wOverQ) * p0 + \
    (t.spacing * wOverQ - 1) * p0.backward

stencil_p0 = Eq(p0.forward, update_p0)

dt = time_axis.step
spacing_map = vel.grid.spacing_map