Skip to content

Commit

Permalink
Merge pull request #88 from trustimaging/format-revamp
Browse files Browse the repository at this point in the history
Format revamp
  • Loading branch information
ccuetom authored Oct 18, 2024
2 parents 5ad3421 + 0a03c54 commit 327d01f
Show file tree
Hide file tree
Showing 4 changed files with 146 additions and 36 deletions.
47 changes: 32 additions & 15 deletions mosaic/file_manipulation/h5.py
Original file line number Diff line number Diff line change
Expand Up @@ -92,8 +92,7 @@ def append(name, obj, group):
append(sub_group_name, obj[index], sub_group)

else:
if name not in group:
_write_dataset(name, obj, group)
_write_dataset(name, obj, group)


def _write_dataset(name, obj, group):
Expand Down Expand Up @@ -128,31 +127,43 @@ def read(obj, lazy=True, filter=None, only=None):
if filter is None:
filter = {}

for key, filter_list in filter.items():
if key in obj:
value = _read_dataset(obj[key], lazy=False)
if value not in filter_list:
raise FilterException
# for key, filter_list in filter.items():
# if key in obj:
# value = _read_dataset(obj[key], lazy=False)
# if value not in filter_list:
# raise FilterException

if obj.attrs.get('is_array'):
data = []
for key in sorted(obj.keys()):
if only is not None and key not in only:
continue
try:
value = read(obj[key], lazy=lazy, filter=filter)
except FilterException:
continue
if filter is not None and key in filter:
try:
value = read(obj[key], lazy=lazy, filter=filter, only=filter[key])
except FilterException:
continue
else:
try:
value = read(obj[key], lazy=lazy, filter=filter)
except FilterException:
continue
data.append(value)
else:
data = {}
for key in obj.keys():
if only is not None and key not in only:
continue
try:
value = read(obj[key], lazy=lazy, filter=filter)
except FilterException:
continue
if filter is not None and key in filter:
try:
value = read(obj[key], lazy=lazy, filter=filter, only=filter[key])
except FilterException:
continue
else:
try:
value = read(obj[key], lazy=lazy, filter=filter)
except FilterException:
continue
data[key] = value

return data
Expand Down Expand Up @@ -306,6 +317,12 @@ def close(self):
self._file.close()

def load(self, lazy=True, filter=None, only=None):
if filter is not None:
for k, v in filter.items():
if not isinstance(v, list):
v = [v]
v = [str(v_) if not isinstance(v_, str) else v_ for v_ in v]
filter[k] = v
group = self._file['/']
description = read(group, lazy=lazy, filter=filter, only=only)
return Struct(description)
Expand Down
55 changes: 43 additions & 12 deletions stride/problem/acquisitions.py
Original file line number Diff line number Diff line change
Expand Up @@ -4,6 +4,7 @@
from collections import OrderedDict
from cached_property import cached_property

import mosaic.types
from mosaic.file_manipulation import h5

from .data import Traces
Expand Down Expand Up @@ -392,15 +393,15 @@ def append_observed(self, *args, **kwargs):
if h5.file_exists(*args, **kwargs):
description = {
'shots': {
'shots_%08d' % self.id: self.__get_desc__(**kwargs)
str(self.id): self.__get_desc__(**kwargs)
}
}

with h5.HDF5(*args, **kwargs, mode='a') as file:
file.append(description)

else:
self._acquisitions.dump(*args, **kwargs)
self._acquisitions.dump(*args, shot_ids=[self.id], **kwargs)

def __get_desc__(self, **kwargs):
description = {
Expand Down Expand Up @@ -1172,17 +1173,40 @@ def sub_problem(self, shot, sub_problem):
return sub_acquisitions

def __get_desc__(self, **kwargs):
description = {
'num_shots': self.num_shots,
'shots': [],
'sequences': [],
}
legacy = kwargs.pop('legacy', False)
shot_ids = kwargs.pop('shot_ids', None)

if legacy:
mosaic.logger().warn('Loading legacy Acquisitions file...')

description = {
'num_shots': self.num_shots,
'shots': [],
'sequences': [],
}

for shot in self.shots:
description['shots'].append(shot.__get_desc__())
for shot in self.shots:
if shot_ids is not None and shot.id not in shot_ids:
continue
description['shots'].append(shot.__get_desc__())

for sequence in self.sequences:
description['sequences'].append(sequence.__get_desc__())
for sequence in self.sequences:
description['sequences'].append(sequence.__get_desc__())

else:
description = {
'num_shots': self.num_shots,
'shots': {},
'sequences': {},
}

for shot in self.shots:
if shot_ids is not None and shot.id not in shot_ids:
continue
description['shots'][str(shot.id)] = shot.__get_desc__()

for sequence in self.sequences:
description['sequences'][str(sequence.id)] = sequence.__get_desc__()

return description

Expand All @@ -1191,6 +1215,9 @@ def __set_desc__(self, description):
shots = description.shots
else:
shots = description.values()
if isinstance(shots, mosaic.types.Struct):
shots = shots.values()

for shot_desc in shots:
if shot_desc.id not in self._shots:
shot = Shot(shot_desc.id,
Expand All @@ -1202,7 +1229,11 @@ def __set_desc__(self, description):
shot.__set_desc__(shot_desc)

if 'sequences' in description:
for seq_desc in description.sequences:
sequences = description.sequences
if isinstance(sequences, mosaic.types.Struct):
sequences = sequences.values()

for seq_desc in sequences:
if seq_desc.id not in self._sequences:
sequence = Sequence(seq_desc.id,
geometry=self._geometry,
Expand Down
4 changes: 4 additions & 0 deletions stride/problem/data.py
Original file line number Diff line number Diff line change
Expand Up @@ -1614,6 +1614,10 @@ def _resample(self, factor, new_num, **kwargs):

if self.allocated:
data = resampy.resample(self.data, sr_orig, sr_new, axis=1) # resample
if data.shape[-1] < new_num:
data = np.pad(data, ((0, 0), (0, new_num-data.shape[-1])), mode='constant', constant_values=0)
elif data.shape[-1] > new_num:
data = data[:, :new_num]
new_traces = Traces(name=self.name, grid=self.grid, transducer_ids=self._transducer_ids, data=data)
else:
new_traces = Traces(name=self.name, grid=self.grid, transducer_ids=self._transducer_ids)
Expand Down
76 changes: 67 additions & 9 deletions stride/problem/geometry.py
Original file line number Diff line number Diff line change
@@ -1,6 +1,7 @@
import numpy as np
from collections import OrderedDict

import mosaic.types
from .base import GriddedSaved, ProblemBase
from .. import plotting
from ..utils import geometries
Expand Down Expand Up @@ -91,10 +92,16 @@ def __get_desc__(self, **kwargs):
def __set_desc__(self, description, transducers=None):
self.id = description.id
self.transducer = transducers.get(description.transducer_id)
self.coordinates = description.coordinates.load()
if hasattr(description.coordinates, 'load'):
self.coordinates = description.coordinates.load()
else:
self.coordinates = description.coordinates

if 'orientation' in description:
self.orientation = description.orientation.load()
if hasattr(description.orientation, 'load'):
self.orientation = description.orientation.load()
else:
self.orientation = description.orientation


class Geometry(ProblemBase):
Expand Down Expand Up @@ -393,18 +400,69 @@ def sub_problem(self, shot, sub_problem):
return sub_geometry

def __get_desc__(self, **kwargs):
description = {
'num_locations': self.num_locations,
'locations': [],
}
legacy = kwargs.pop('legacy', False)

for location_id, location in self._locations.items():
description['locations'].append(location.__get_desc__())
if legacy:
mosaic.logger().warn('Loading legacy Geometry file...')

description = {
'num_locations': self.num_locations,
'locations': [],
}

for location_id, location in self._locations.items():
description['locations'].append(location.__get_desc__())

else:
description = {
'num_locations': self.num_locations,
'locations': {},
}

coordinates = None
orientation = None

index = 0
for location_id, location in self._locations.items():
location_desc = location.__get_desc__()
if 'coordinates' in location_desc:
if coordinates is None:
coordinates = np.zeros((self.num_locations, self.space.dim), dtype=np.float32)
coordinates[index, :] = location_desc['coordinates']
location_desc['coordinates'] = index
if 'orientation' in location_desc:
if orientation is None:
orientation = np.zeros((self.num_locations, self.space.dim), dtype=np.float32)
orientation[index, :] = location_desc['orientation']
location_desc['orientation'] = index

description['locations'][str(location_id)] = location_desc
index += 1

if coordinates is not None:
description['coordinates'] = coordinates

if orientation is not None:
description['orientation'] = orientation

return description

def __set_desc__(self, description):
for location_desc in description.locations:
locations = description.locations
if isinstance(locations, mosaic.types.Struct):
locations = locations.values()

if 'coordinates' in description:
for location_desc in locations:
idx = location_desc.coordinates
location_desc.coordinates = description.coordinates[idx, :]

if 'orientation' in description:
for location_desc in locations:
idx = location_desc.orientation
location_desc.orientation = description.orientation[idx, :]

for location_desc in locations:
if location_desc.id not in self.location_ids:
instance = TransducerLocation(location_desc.id)
self.add_location(instance)
Expand Down

0 comments on commit 327d01f

Please sign in to comment.