Skip to content

Commit

Permalink
fix some lint
Browse files Browse the repository at this point in the history
  • Loading branch information
jshipton committed Aug 24, 2023
1 parent cd0232f commit af283f9
Show file tree
Hide file tree
Showing 2 changed files with 12 additions and 16 deletions.
14 changes: 5 additions & 9 deletions gusto/moving_mesh/monge_ampere.py
Original file line number Diff line number Diff line change
@@ -1,6 +1,7 @@
from firedrake import *
import numpy as np


class OptimalTransportMeshGenerator(object):

def __init__(self, mesh_in, monitor, initial_tol=1.e-4, tol=1.e-2, pre_meshgen_callback=None, post_meshgen_callback=None):
Expand Down Expand Up @@ -42,7 +43,7 @@ def setup(self):
# get mesh area
self.total_area = assemble(Constant(1.0)*dx(self.mesh))

### SET UP FUNCTIONS ###
# set up functions
P2 = FunctionSpace(self.mesh, "Q" if quads else "P", 2)
TensorP2 = TensorFunctionSpace(self.mesh, "Q" if quads else "P", 2)
MixedSpace = P2*TensorP2
Expand All @@ -65,7 +66,7 @@ def setup(self):
self.m = Function(P1)
self.theta = Constant(0.0)

### DEFINE MESH EQUATIONS ###
# define mesh equations
v, tau = TestFunctions(MixedSpace)

if self.meshtype == "plane":
Expand Down Expand Up @@ -101,7 +102,7 @@ def setup(self):
elif self.meshtype == "sphere":
self.gradphi_cts2 = Function(W_cts) # extra, as gradphi_cts not necessarily tangential

### SET UP INITIAL SIGMA ### (needed on sphere)
# set up initial sigma (needed on sphere)
sigma_ = TrialFunction(TensorP2)
tau_ = TestFunction(TensorP2)
sigma_ini = Function(TensorP2)
Expand All @@ -116,7 +117,7 @@ def setup(self):

self.phisigma.sub(1).assign(sigma_ini)

### SOLVER OPTIONS FOR MESH GENERATION ###
# solver options for mesh generation
phi__, sigma__ = TrialFunctions(MixedSpace)
v__, tau__ = TestFunctions(MixedSpace)

Expand Down Expand Up @@ -197,14 +198,11 @@ def update_mxtheta(self, cursol):
self.own_output_coords.interpolate(Constant(self.R)*self.x)
self.output_coords.dat.data[:] = self.own_output_coords.dat.data_ro[:]
self.mesh_in.coordinates.assign(self.output_coords)
#self.mesh_in.coordinates.assign(self.x)
self.initialise_fn()
#self.monitor.mesh.coordinates.dat.data[:] = self.mesh.coordinates.dat.data_ro[:]
self.monitor.mesh.coordinates.dat.data[:] = self.own_output_coords.dat.data_ro[:]
self.monitor.update_monitor()
self.m.dat.data[:] = self.monitor.m.dat.data_ro[:]
else:
# self.mesh.coordinates.assign(self.x)
self.own_output_coords.interpolate(Constant(self.R)*self.x)
self.x_new.dat.data[:] = self.own_output_coords.dat.data_ro[:]
self.monitor.get_monitor_on_new_mesh(self.monitor.m, self.x_old, self.x_new)
Expand Down Expand Up @@ -264,8 +262,6 @@ def get_new_mesh(self):

# Move data from internal mesh to output mesh.
self.own_output_coords.interpolate(Constant(self.R)*self.x)
# self.mesh.coordinates.assign(self.x)
# self.output_coords.dat.data[:] = self.mesh.coordinates.dat.data_ro[:]
self.output_coords.dat.data[:] = self.own_output_coords.dat.data_ro[:]
return self.output_coords

Expand Down
14 changes: 7 additions & 7 deletions gusto/moving_mesh/monitor.py
Original file line number Diff line number Diff line change
@@ -1,10 +1,10 @@
from firedrake import *
import math
import numpy as np
from mpi4py import MPI

__all__ = ["MonitorFunction"]


class MonitorFunction(object):

def __init__(self, f_name, adapt_to="function", avg_weight=0.5,
Expand All @@ -30,7 +30,7 @@ def setup(self, state_fields):
self.f = Function(FunctionSpace(self.mesh, f.ufl_element())) # make "own copy" of f on internal mesh
self.user_f = f

### SET UP FUNCTION SPACES ###
# Set up function spaces
P1 = FunctionSpace(self.mesh, "Q" if quads else "P", 1) # for representing m
DP1 = FunctionSpace(self.mesh, "DQ" if quads else "DP", 1) # for advection
VectorP1 = VectorFunctionSpace(self.mesh, "Q" if quads else "P", 1)
Expand All @@ -43,15 +43,15 @@ def setup(self, state_fields):

# get mesh area
self.total_area = assemble(Constant(1.0)*dx(self.mesh))

self.m = Function(P1)
self.m_prereg = Function(P1)
self.m_old = Function(P1)
self.m_dg = Function(DP1)
self.dm = Function(DP1)
self.m_int_form = self.m_prereg*dx

### DEFINE MONITOR FUNCTION IN TERMS OF q ###
# define monitor function in terms of q
if False: # for plane
v_ones = as_vector(np.ones(2))
else:
Expand All @@ -72,7 +72,7 @@ def setup(self, state_fields):
v_tp1 = TestFunction(TensorP1)
self.a_tp1_lumped = inner(v_tp1, t_ones)*dx
self.L_tp1 = inner(v_tp1, grad(self.gradq))*dx

# Define forms for lumped project of monitor function into P1
v_p1 = TestFunction(P1)
self.a_p1_lumped = v_p1*dx
Expand Down Expand Up @@ -129,7 +129,7 @@ def update_monitor(self):
m_min = self.m.comm.allreduce(self.m.dat.data_ro.min(), op=MPI.MIN)
self.m.dat.data[:] /= m_min

mmax_pre = max(self.m.dat.data)/min(self.m.dat.data)
# mmax_pre = max(self.m.dat.data)/min(self.m.dat.data)

def get_monitor_on_new_mesh(self, m, x_old, x_new):
# We have the function m on old mesh: m_old. We want to represent
Expand All @@ -150,4 +150,4 @@ def get_monitor_on_new_mesh(self, m, x_old, x_new):
self.limiter.apply(self.m_dg)

project(self.m_dg, self.m) # project discontinuous m back into CG
mmax_post = max(self.m.dat.data)/min(self.m.dat.data)
# mmax_post = max(self.m.dat.data)/min(self.m.dat.data)

0 comments on commit af283f9

Please sign in to comment.