Skip to content

Commit

Permalink
Black updates to the rest of datamodel
Browse files Browse the repository at this point in the history
  • Loading branch information
@rieder
rieder committed Sep 17, 2024
1 parent ed2dede commit 73724b7
Show file tree
Hide file tree
Showing 10 changed files with 2,781 additions and 2,129 deletions.
1,204 changes: 687 additions & 517 deletions src/amuse/datamodel/base.py
View file

Large diffs are not rendered by default.

1,150 changes: 651 additions & 499 deletions src/amuse/datamodel/grids.py
View file

Large diffs are not rendered by default.

321 changes: 179 additions & 142 deletions src/amuse/datamodel/indexing.py
View file

Large diffs are not rendered by default.

467 changes: 245 additions & 222 deletions src/amuse/datamodel/memory_storage.py
View file

Large diffs are not rendered by default.

416 changes: 225 additions & 191 deletions src/amuse/datamodel/parameters.py
View file

Large diffs are not rendered by default.

807 changes: 492 additions & 315 deletions src/amuse/datamodel/particle_attributes.py
View file

Large diffs are not rendered by default.

13 changes: 2 additions & 11 deletions src/amuse/datamodel/particles.py
View file
Original file line number Diff line number Diff line change
Expand Up @@ -334,7 +334,6 @@ def to_string(self, attributes_to_show=None, split_at=20):

rows = []
for i in range(len(columns[0])):

row = [x[i] for x in columns]
rows.append(row)

Expand Down Expand Up @@ -1175,7 +1174,7 @@ def __init__(
keys_generator=None,
particles=None,
is_working_copy=True,
**attributes
**attributes,
):
AbstractParticleSet.__init__(self)

Expand Down Expand Up @@ -1535,15 +1534,13 @@ def __call__(self, *list_arguments, **keyword_arguments):


class DerivedSupersetAttribute(DerivedAttribute):

def __init__(self, name):
self.name = name

def get_values_for_entities(self, superset):
result = None
offset = 0
for subset in superset._private.particle_sets:

subset_result = getattr(subset, self.name)
if hasattr(subset_result, "__call__"):
if len(subset) > 0:
Expand Down Expand Up @@ -3040,7 +3037,6 @@ def get_valid_particles_mask(self):
return self._private.particles.get_valid_particles_mask()

def __getitem__(self, index):

keys = self.get_all_keys_in_store()[index]

if keys is ma.masked:
Expand Down Expand Up @@ -3217,7 +3213,6 @@ def get_valid_particles_mask(self):
return self._private.particles.get_valid_particles_mask()

def __getitem__(self, index):

keys = self.get_all_keys_in_store()[index]

if keys is ma.masked:
Expand Down Expand Up @@ -3441,7 +3436,6 @@ def shallow_copy(self):
return self._factory(copiedParticles)

def get_values_in_store(self, indices, attributes):

direct_attributes = []
needs_transformation = False
for attribute in attributes:
Expand Down Expand Up @@ -3590,7 +3584,6 @@ def get_subset(self, name):


class ParticleInformationChannel(object):

def __init__(
self, from_particles, to_particles, attributes=None, target_names=None
):
Expand Down Expand Up @@ -3663,7 +3656,6 @@ def copy_attributes(self, attributes, target_names=None):
def copy_attributes_async(
self, attributes, target_names=None, async_get=True, async_set=False
):

if target_names is None:
target_names = attributes

Expand Down Expand Up @@ -3882,7 +3874,6 @@ def get_valid_particles_mask(self):
return self._private.particles.get_valid_particles_mask()

def __getitem__(self, index):

keys = self.get_all_keys_in_store()[index]

if keys is ma.masked:
Expand Down Expand Up @@ -4054,7 +4045,7 @@ def __init__(
particles_set=None,
set_index=None,
set_version=-1,
**keyword_arguments
**keyword_arguments,
):
if particles_set is None:
if key == None:
Expand Down
189 changes: 105 additions & 84 deletions src/amuse/datamodel/staggeredgrid.py
View file
Original file line number Diff line number Diff line change
Expand Up @@ -6,9 +6,9 @@


class StaggeredGrid(object):
def __init__(self, elements, nodes = None, get_corners=None):
self.elements=elements
self.nodes=nodes
def __init__(self, elements, nodes=None, get_corners=None):
self.elements = elements
self.nodes = nodes
self._get_corners_func = get_corners
self.get_corners()

Expand All @@ -29,104 +29,109 @@ def get_corners(self):
elements = self.elements
nodes = self.nodes

#if a replacing get_corners method is defined for this grid call it and return
# if a replacing get_corners method is defined for this grid call it and return
if not (self._get_corners_func is None):
corners = self._get_corners_func()
if not hasattr(elements, '_cell_corners'):
if not hasattr(elements, "_cell_corners"):
object.__setattr__(elements, "_cell_corners", corners)
self.corners = corners
return corners
return corners

#element and node grids should be of the same type
# element and node grids should be of the same type
if type(elements) != type(nodes):
raise Exception("elements and nodes grids should be of the same type")

#get set of participating axes names
# get set of participating axes names
dims = elements.number_of_dimensions()
if dims > 2:
raise Exception("Staggered grids with more than 2 dimensions are currently not supported")
#this is mainly because we can only access the coordinates in
#a dimensions by their attribute names 'lat' and 'lon'
#if the _axes_names were defined properly for in-code stored grids I could do something like
raise Exception(
"Staggered grids with more than 2 dimensions are currently not supported"
)
# this is mainly because we can only access the coordinates in
# a dimensions by their attribute names 'lat' and 'lon'
# if the _axes_names were defined properly for in-code stored grids I could do something like
# xpos=to_quantity(getattr(nodes,self._axes_names[0])), but currently that's not possible right now

#structured or unstructured
# structured or unstructured
if type(elements) is StructuredGrid:

#structured grid assumes that the nodes are at the cell corners
#of the elements grid and therefore the nodes grid should be
#exactly 1 gridpoint larger in each dimension to fully encapsulate
#the element grid
# structured grid assumes that the nodes are at the cell corners
# of the elements grid and therefore the nodes grid should be
# exactly 1 gridpoint larger in each dimension to fully encapsulate
# the element grid
if len(elements.shape) != len(nodes.shape):
raise Exception("elements and nodes should have the same number of the dimensions")
raise Exception(
"elements and nodes should have the same number of the dimensions"
)

for i in range(len(elements.shape)):
if elements.shape[i]+1 != nodes.shape[i]:
raise Exception("nodes grid should be exactly 1 grid point larger than element grid in each dimension")
if elements.shape[i] + 1 != nodes.shape[i]:
raise Exception(
"nodes grid should be exactly 1 grid point larger than element grid in each dimension"
)

corners = numpy.zeros([dims] + list(nodes.shape), dtype=numpy.double)

#use node positions as corner positions
# use node positions as corner positions
corners[0] = nodes.lon.value_in(units.rad)
corners[1] = nodes.lat.value_in(units.rad)

elif type(elements) == UnstructuredGrid:
#the following bit of code tries to access 'n0', 'n1' up to 'n9' of the elements grid
#a cleaner implementation would be to call get_element_nodes() directly, but we can not do that from here
# the following bit of code tries to access 'n0', 'n1' up to 'n9' of the elements grid
# a cleaner implementation would be to call get_element_nodes() directly, but we can not do that from here
attributes = elements.all_attributes()
max_corners = 10
corner_indices = []
for i in range(max_corners):
node = 'n' + str(i)
node = "n" + str(i)
if node in attributes:
corner_indices.append(getattr(elements, node))

self.num_corners = num_corners = len(corner_indices)
object.__setattr__(elements,"_num_corners", num_corners)
object.__setattr__(elements, "_num_corners", num_corners)
self.corner_indices = corner_indices
size = elements.size
self.inverse_mapping = inverse_mapping = [[] for i in range(nodes.size)]

#only 2 dimensions supported currently
corners = numpy.zeros((2, size*num_corners), dtype=numpy.double)
# only 2 dimensions supported currently
corners = numpy.zeros((2, size * num_corners), dtype=numpy.double)
node_lon = nodes.lon.value_in(units.rad)
node_lat = nodes.lat.value_in(units.rad)
for i in range(size):
for d in range(num_corners):
n = corner_indices[d][i]
inverse_mapping[n].append(i)
nlon = node_lon[n]
corners[0][i*num_corners+d] = nlon
corners[0][i * num_corners + d] = nlon
nlat = node_lat[n]
corners[1][i*num_corners+d] = nlat

else:
raise Exception("unknown grid type for elements: should be either StructuredGrid or UnstructuredGrid")

corners[1][i * num_corners + d] = nlat

else:
raise Exception(
"unknown grid type for elements: should be either StructuredGrid or UnstructuredGrid"
)

if not hasattr(elements, '_cell_corners'):
if not hasattr(elements, "_cell_corners"):
object.__setattr__(elements, "_cell_corners", corners)
self.corners = corners
return corners
return corners

def map_elements_to_nodes_structured_larger(self, elements, nodes, elem_values):
node_values = numpy.zeros(self.nodes.shape, dtype=numpy.float64)
node_values[1:,1:] = elem_values[:,:]
node_values[0,:] = 0.0
#assume the grid is cyclic east-west
node_values[1:,0] = elem_values[:,-1]
node_values[1:, 1:] = elem_values[:, :]
node_values[0, :] = 0.0
# assume the grid is cyclic east-west
node_values[1:, 0] = elem_values[:, -1]
return node_values

def map_elements_to_nodes_structured_same_size(self, elements, nodes, elem_values):
node_values = numpy.zeros(self.nodes.shape, dtype=numpy.float64)
node_values = elem_values[:]
return node_values

def map_elements_to_nodes(self, element_values):
#currently very rough remapping schemes, more sophisticated methods will be added later
# currently very rough remapping schemes, more sophisticated methods will be added later

if not hasattr(self, 'corners'):
if not hasattr(self, "corners"):
self.get_corners()
elements = self.elements
nodes = self.nodes
Expand All @@ -135,52 +140,61 @@ def map_elements_to_nodes(self, element_values):

if type(elements) is StructuredGrid:
if len(elements.shape) != len(nodes.shape):
raise Exception("elements and nodes should have the same number of the dimensions")

if numpy.all([s1+1==s2 for s1,s2 in zip(elements.shape,nodes.shape)]):
return self.map_elements_to_nodes_structured_larger(elements, nodes, element_values)
if numpy.all([s1==s2 for s1,s2 in zip(elements.shape,nodes.shape)]):
return self.map_elements_to_nodes_structured_same_size(elements, nodes, element_values)
raise Exception(
"elements and nodes should have the same number of the dimensions"
)

if numpy.all([s1 + 1 == s2 for s1, s2 in zip(elements.shape, nodes.shape)]):
return self.map_elements_to_nodes_structured_larger(
elements, nodes, element_values
)
if numpy.all([s1 == s2 for s1, s2 in zip(elements.shape, nodes.shape)]):
return self.map_elements_to_nodes_structured_same_size(
elements, nodes, element_values
)
else:
raise Exception("nodes grid should have either exactly same shape or 1 grid point larger than element grid in each dimension")
raise Exception(
"nodes grid should have either exactly same shape or 1 grid point larger than element grid in each dimension"
)

elif type(elements) == UnstructuredGrid:
if (len(element_values) != self.elements.size):
raise Exception("number of values passed does not match size of elements grid")
#do a simple average value of the elements around the node
if len(element_values) != self.elements.size:
raise Exception(
"number of values passed does not match size of elements grid"
)
# do a simple average value of the elements around the node
node_values = numpy.zeros(self.nodes.size, dtype=numpy.float64)
for i in range(len(node_values)):
num_neighbors = len(self.inverse_mapping[i])
value = 0.0
#add value of neighboring element
# add value of neighboring element
for neighbor in self.inverse_mapping[i]:
value += element_values[neighbor]
#store result
node_values[i] = value / (1.0*num_neighbors)
# store result
node_values[i] = value / (1.0 * num_neighbors)
else:
raise Exception("unknown grid type for elements: should be either StructuredGrid or UnstructuredGrid")
raise Exception(
"unknown grid type for elements: should be either StructuredGrid or UnstructuredGrid"
)
return node_values


#this method is for structured grids where the nodes grid is exactly 1 grid point larger in each dimension
# this method is for structured grids where the nodes grid is exactly 1 grid point larger in each dimension
def map_nodes_to_elements_structured_larger(self, elements, nodes, node_values):
#do simple translation/shift of the values from the north-east corners of each grid cell to the cell centers
# do simple translation/shift of the values from the north-east corners of each grid cell to the cell centers
elem_values = numpy.zeros(self.elements.shape, dtype=numpy.float64)
elem_values = node_values[1:,1:]
elem_values = node_values[1:, 1:]
return elem_values

#this method is for structured grids where the nodes grid is of the same size as the elements grid, if so
#the grid is assumed to be cyclic east-west
# this method is for structured grids where the nodes grid is of the same size as the elements grid, if so
# the grid is assumed to be cyclic east-west
def map_nodes_to_elements_structured_same_size(self, elements, nodes, node_values):
#do simple translation/shift of the values from the north-east corners of each grid cell to the cell centers
# do simple translation/shift of the values from the north-east corners of each grid cell to the cell centers
elem_values = numpy.zeros(self.elements.shape, dtype=numpy.float64)
elem_values = node_values.flatten()
return elem_values


def map_nodes_to_elements(self, node_values):

if not hasattr(self, 'corners'):
if not hasattr(self, "corners"):
self.get_corners()
elements = self.elements
nodes = self.nodes
Expand All @@ -189,33 +203,40 @@ def map_nodes_to_elements(self, node_values):

if type(elements) is StructuredGrid:
if len(elements.shape) != len(nodes.shape):
raise Exception("elements and nodes should have the same number of the dimensions")
if numpy.all([s1+1==s2 for s1,s2 in zip(elements.shape,nodes.shape)]):
return self.map_nodes_to_elements_structured_larger(elements, nodes, node_values)
if numpy.all([s1==s2 for s1,s2 in zip(elements.shape,nodes.shape)]):
return self.map_nodes_to_elements_structured_same_size(elements, nodes, node_values)
raise Exception(
"elements and nodes should have the same number of the dimensions"
)
if numpy.all([s1 + 1 == s2 for s1, s2 in zip(elements.shape, nodes.shape)]):
return self.map_nodes_to_elements_structured_larger(
elements, nodes, node_values
)
if numpy.all([s1 == s2 for s1, s2 in zip(elements.shape, nodes.shape)]):
return self.map_nodes_to_elements_structured_same_size(
elements, nodes, node_values
)
else:
raise Exception("nodes grid should have either exactly same shape or 1 grid point larger than element grid in each dimension")

raise Exception(
"nodes grid should have either exactly same shape or 1 grid point larger than element grid in each dimension"
)

elif type(elements) == UnstructuredGrid:
if (len(node_values) != self.nodes.size):
raise Exception("number of values passed does not match size of nodes grid")
if len(node_values) != self.nodes.size:
raise Exception(
"number of values passed does not match size of nodes grid"
)

elem_values = numpy.zeros(self.elements.size, dtype=numpy.float64)

#do a simple average value of the nodes around the element
# do a simple average value of the nodes around the element
for i in range(len(elem_values)):
value = 0.0
for c in range(self.num_corners):
index = self.corner_indices[c][i]
value += node_values[index]
elem_values[i] = value / (1.0*self.num_corners)
value += node_values[index]
elem_values[i] = value / (1.0 * self.num_corners)

else:
raise Exception("unknown grid type for elements: should be either StructuredGrid or UnstructuredGrid")
raise Exception(
"unknown grid type for elements: should be either StructuredGrid or UnstructuredGrid"
)
return elem_values




Loading

0 comments on commit 73724b7

Please sign in to comment.