diff --git a/.gitignore b/.gitignore index 132455c..d2a0424 100644 --- a/.gitignore +++ b/.gitignore @@ -13,3 +13,4 @@ dist/ *.rst *.png *.mp4 +*.log diff --git a/paicos/__init__.py b/paicos/__init__.py index 04a7ec8..47635a4 100644 --- a/paicos/__init__.py +++ b/paicos/__init__.py @@ -44,6 +44,7 @@ from .image_creators.nested_projector import NestedProjector from .image_creators.tree_projector import TreeProjector from .image_creators.slicer import Slicer +from .image_creators.image_creator_actions import Actions # Histograms from .histograms.histogram import Histogram diff --git a/paicos/cython/sph_projectors.pyx b/paicos/cython/sph_projectors.pyx index 9c10bb5..d0f9350 100644 --- a/paicos/cython/sph_projectors.pyx +++ b/paicos/cython/sph_projectors.pyx @@ -51,9 +51,10 @@ def project_image(real_t[:] xvec, real_t[:] yvec, real_t[:] variable, cdef int Np = xvec.shape[0] # Shape of projection array - cdef int ny = (sidelength_y/sidelength_x * nx) + cdef int ny = round(sidelength_y/sidelength_x * nx) msg = '(sidelength_y/sidelength_x * nx) needs to be an integer' - assert (sidelength_y/sidelength_x * nx) == ny, msg + + assert abs((sidelength_y/sidelength_x * nx) - ny) < 1e-8, msg # Create projection array cdef real_t[:, :] projection = np.zeros((nx, ny), dtype=np.float64) @@ -67,7 +68,7 @@ def project_image(real_t[:] xvec, real_t[:] yvec, real_t[:] variable, cdef real_t x, y, norm cdef real_t x0, y0 - assert sidelength_x/nx == sidelength_y/ny + # assert sidelength_x/nx == sidelength_y/ny # Lower left corner of image in Arepo coordinates x0 = xc - sidelength_x / 2.0 @@ -159,9 +160,10 @@ def project_image_omp(real_t[:] xvec, real_t[:] yvec, real_t[:] variable, cdef int Np = xvec.shape[0] # Shape of projection array - cdef int ny = (sidelength_y/sidelength_x * nx) + cdef int ny = round(sidelength_y/sidelength_x * nx) msg = '(sidelength_y/sidelength_x * nx) needs to be an integer' - assert (sidelength_y/sidelength_x * nx) == ny, msg + + assert abs((sidelength_y/sidelength_x * nx) - ny) < 1e-8, msg # Loop integers and other variables cdef int ip, ix, iy, ih, threadnum @@ -172,7 +174,7 @@ def project_image_omp(real_t[:] xvec, real_t[:] yvec, real_t[:] variable, cdef real_t x, y, norm cdef real_t x0, y0 - assert sidelength_x/nx == sidelength_y/ny + # assert sidelength_x/nx == sidelength_y/ny # Lower left corner of image in Arepo coordinates x0 = xc - sidelength_x/2.0 @@ -282,9 +284,10 @@ def project_oriented_image(real_t[:] xvec, real_t[:] yvec, real_t[:] zvec, real_ cdef int Np = xvec.shape[0] # Shape of projection array - cdef int ny = (sidelength_y/sidelength_x * nx) + cdef int ny = round(sidelength_y/sidelength_x * nx) msg = '(sidelength_y/sidelength_x * nx) needs to be an integer' - assert (sidelength_y/sidelength_x * nx) == ny, msg + + assert abs((sidelength_y/sidelength_x * nx) - ny) < 1e-8, msg # Create projection array cdef real_t[:, :] projection = np.zeros((nx, ny), dtype=np.float64) @@ -298,7 +301,7 @@ def project_oriented_image(real_t[:] xvec, real_t[:] yvec, real_t[:] zvec, real_ cdef real_t x, y, norm cdef real_t cen_x, cen_y, cen_z - assert sidelength_x/nx == sidelength_y/ny + # assert sidelength_x/nx == sidelength_y/ny for ip in range(Np): # Centered coordinates @@ -396,9 +399,10 @@ def project_oriented_image_omp(real_t[:] xvec, real_t[:] yvec, real_t[:] zvec, r cdef int Np = xvec.shape[0] # Shape of projection array - cdef int ny = (sidelength_y/sidelength_x * nx) + cdef int ny = round(sidelength_y/sidelength_x * nx) msg = '(sidelength_y/sidelength_x * nx) needs to be an integer' - assert (sidelength_y/sidelength_x * nx) == ny, msg + + assert abs((sidelength_y/sidelength_x * nx) - ny) < 1e-8, msg # Loop integers and other variables cdef int ip, ix, iy, ih, threadnum @@ -409,7 +413,7 @@ def project_oriented_image_omp(real_t[:] xvec, real_t[:] yvec, real_t[:] zvec, r cdef real_t x, y, norm cdef real_t cen_x, cen_y, cen_z - assert sidelength_x/nx == sidelength_y/ny + #assert sidelength_x/nx == sidelength_y/ny # Create projection arrays cdef real_t[:, :, :] tmp_var = np.zeros((nx, ny, numthreads), dtype=np.float64) diff --git a/paicos/image_creators/gpu_ray_projector.py b/paicos/image_creators/gpu_ray_projector.py index 603bed2..7748584 100644 --- a/paicos/image_creators/gpu_ray_projector.py +++ b/paicos/image_creators/gpu_ray_projector.py @@ -257,17 +257,22 @@ def _send_small_data_to_gpu(self): self.gpu_variables['widths'] = cp.array(self.widths) self.gpu_variables['center'] = cp.array(self.center) - def _gpu_project(self, variable_str): + def _gpu_project(self, variable_str, additive): """ Private method for projecting using cuda code """ - rotation_matrix = self.gpu_variables['rotation_matrix'] - widths = self.gpu_variables['widths'] - center = self.gpu_variables['center'] - hsml = self.gpu_variables['hsml'] + gpu_vars = self.gpu_variables + rotation_matrix = gpu_vars['rotation_matrix'] + widths = gpu_vars['widths'] + center = gpu_vars['center'] + hsml = gpu_vars['hsml'] nx = self.npix_width ny = self.npix_height - variable = self.gpu_variables[variable_str] + + if additive: + variable = gpu_vars[variable_str] / gpu_vars[f'{self.parttype}_Volume'] + else: + variable = gpu_vars[variable_str] tree_scale_factor = self.tree.conversion_factor tree_offsets = self.tree.off_sets @@ -286,29 +291,7 @@ def _gpu_project(self, variable_str): projection = cp.asnumpy(image) return projection - def project_variable(self, variable, additive=False): - """ - projects a given variable onto a 2D plane. - - Parameters - ---------- - variable : str, function, numpy array - variable, it can be passed as string or an array - - Returns - ------- - numpy array - The image of the projected variable - """ - - # This calls _do_region_selection if resolution, Orientation, - # widths or center changed - self._check_if_properties_changed() - - if additive: - err_msg = "GPU ray tracer does not yet support additive=True" - raise RuntimeError(err_msg) - + def _send_variable_to_gpu(self, variable, gpu_key='projection_variable'): if isinstance(variable, str): variable_str = str(variable) err_msg = 'projector uses a different parttype' @@ -323,9 +306,10 @@ def project_variable(self, variable, additive=False): # Select same part of array that the projector has selected if self.do_pre_selection: + # TODO: Check that this is not applied more than once... variable = variable[self.index] - if variable_str in self.gpu_variables and variable_str != 'projection_variable': + if variable_str in self.gpu_variables and variable_str != gpu_key: pass else: # Send variable to gpu @@ -338,8 +322,38 @@ def project_variable(self, variable, additive=False): self.gpu_variables[variable_str] = self.gpu_variables[variable_str][ self.tree.sort_index] + if isinstance(variable, units.PaicosQuantity): + unit_quantity = variable.unit_quantity + else: + unit_quantity = None + + return variable_str, unit_quantity + + def project_variable(self, variable, additive=False): + """ + projects a given variable onto a 2D plane. + + Parameters + ---------- + variable : str, function, numpy array + variable, it can be passed as string or an array + + Returns + ------- + numpy array + The image of the projected variable + """ + + # This calls _do_region_selection if resolution, Orientation, + # widths or center changed + self._check_if_properties_changed() + + variable_str, unit_quantity = self._send_variable_to_gpu(variable) + if additive: + _, vol_unit_quantity = self._send_variable_to_gpu(f'{self.parttype}_Volume') + # Do the projection - projection = self._gpu_project(variable_str) + projection = self._gpu_project(variable_str, additive) # Transpose projection = projection.T @@ -347,16 +361,30 @@ def project_variable(self, variable, additive=False): assert projection.shape[0] == self.npix_height assert projection.shape[1] == self.npix_width - if isinstance(variable, units.PaicosQuantity): + if unit_quantity is not None: unit_length = self.snap['0_Coordinates'].uq - projection = projection * variable.unit_quantity * unit_length + projection = projection * unit_quantity * unit_length + if additive: + projection = projection / vol_unit_quantity - return projection / self.depth + if additive: + return projection + else: + return projection / self.depth def __del__(self): """ Clean up like this? Not sure it is needed... """ - del self.gpu_variables - del self.tree + self.release_gpu_memory() + + def release_gpu_memory(self): + if hasattr(self, 'gpu_variables'): + for key in list(self.gpu_variables): + del self.gpu_variables[key] + del self.gpu_variables + if hasattr(self, 'tree'): + self.tree.release_gpu_memory() + del self.tree + cp._default_memory_pool.free_all_blocks() diff --git a/paicos/image_creators/gpu_sph_projector.py b/paicos/image_creators/gpu_sph_projector.py index 83c1e3f..e4db6d1 100644 --- a/paicos/image_creators/gpu_sph_projector.py +++ b/paicos/image_creators/gpu_sph_projector.py @@ -439,7 +439,12 @@ def __del__(self): """ Clean up like this? Not sure it is needed... """ - del self.gpu_variables - # for key in list(self.gpu_variables.keys): - # del self.gpu_variables[key] + self.release_gpu_memory() + + def release_gpu_memory(self): + if hasattr(self, 'gpu_variables'): + for key in list(self.gpu_variables): + del self.gpu_variables[key] + del self.gpu_variables + cp._default_memory_pool.free_all_blocks() diff --git a/paicos/image_creators/image_creator.py b/paicos/image_creators/image_creator.py index 9804b26..27d75e7 100644 --- a/paicos/image_creators/image_creator.py +++ b/paicos/image_creators/image_creator.py @@ -190,6 +190,14 @@ def center(self, value): self._center = np.array(value) self._properties_changed = True + def reset_center(self): + """ + Resets the center of the image creator to the + value it had at initialization. + """ + self.center = self.center - self._diff_center + self._diff_center = np.zeros_like(self._diff_center) + @property def npix(self): """ @@ -365,13 +373,13 @@ def width(self, value): if settings.use_units: assert hasattr(value, 'unit') if self.direction == 'x': - self._widths[1] = value.copy + self._widths[1] = np.copy(value) elif self.direction == 'y': - self._widths[2] = value.copy + self._widths[2] = np.copy(value) elif self.direction == 'z' or self.direction == 'orientation': - self._widths[0] = value.copy + self._widths[0] = np.copy(value) self._properties_changed = True @height.setter @@ -379,13 +387,13 @@ def height(self, value): if settings.use_units: assert hasattr(value, 'unit') if self.direction == 'x': - self._widths[2] = value.copy + self._widths[2] = np.copy(value) elif self.direction == 'y': - self._widths[0] = value.copy + self._widths[0] = np.copy(value) elif self.direction == 'z' or self.direction == 'orientation': - self._widths[1] = value.copy + self._widths[1] = np.copy(value) self._properties_changed = True @depth.setter @@ -393,13 +401,13 @@ def depth(self, value): if settings.use_units: assert hasattr(value, 'unit') if self.direction == 'x': - self._widths[0] = value.copy + self._widths[0] = np.copy(value) elif self.direction == 'y': - self._widths[1] = value.copy + self._widths[1] = np.copy(value) elif self.direction == 'z' or self.direction == 'orientation': - self._widths[2] = value.copy + self._widths[2] = np.copy(value) self._properties_changed = True def double_resolution(self): @@ -485,9 +493,9 @@ def npix_height(self): The number of pixels of an image in the vertical direction. """ if settings.use_units: - return int((self.height / self.width).value * self.npix_width) + return int(round((self.height / self.width).value * self.npix_width)) else: - return int((self.height / self.width) * self.npix_width) + return int(round((self.height / self.width) * self.npix_width)) @property def area(self): diff --git a/paicos/image_creators/image_creator_actions.py b/paicos/image_creators/image_creator_actions.py new file mode 100644 index 0000000..0728994 --- /dev/null +++ b/paicos/image_creators/image_creator_actions.py @@ -0,0 +1,355 @@ +import numpy as np +from ..writers.arepo_image import ImageWriter +from ..readers.paicos_readers import ImageReader +import os + + +def find_angle_subdivisions(angle, max_angle): + kk = 1 + angle = abs(angle) + new_angle = angle / kk + while new_angle > max_angle: + kk += 1 + new_angle = angle / kk + + return kk + + +def find_zoom_sub_divisions(zoom, zoom_max): + kk = 1 + if zoom < 1: + zoom = 1 / zoom + new_zoom = zoom**(1 / kk) + while new_zoom > zoom_max: + kk += 1 + new_zoom = zoom**(1 / kk) + + return kk + + +class Actions: + """ + The purpose of this class is to simplify the steps + that are usually involved in creating animations + and/or interactive widgets. + + This is very much work in progress and subject + to change. Send Thomas an email if you would like + a working example. + """ + def __init__(self, snapnum=None, image_creator=None, dry_run=True): + + self.dry_run = dry_run + + self.first_call = True + if snapnum is None: + if image_creator is not None: + self.image_creator = image_creator + else: + msg = ("self.image_creator not set, you'll " + + "nedd to do this manually or read " + + "a log file using Actions.read_log") + print(msg) + else: + self.make_image_creator(snapnum) + self.first_call = False + + self.logging = False + self.zoom_max = 10000000. + self.max_angle = 360. + + def make_image_creator(self, snapnum): + """ + + TODO: Add example implementation in the docstring of + of this method. + """ + err_msg = ("You need to implement this yourself.\n" + + "Please see the docstring of this method") + raise RuntimeError(err_msg) + + def mylogger(self, string, line_ending='\n', mode='a', command=True): + outfolder = self.outfolder + basename = self.basename + with open(f'{outfolder}/{basename}.log', mode) as f: + f.write(string + line_ending) + + if self.make_waypoints and command: + self.frame_num += 1 + waypoint_folder = f'{self.outfolder}/{self.basename}_waypoints' + self._make_waypoint(waypoint_folder, self.basename, self.frame_num) + + def create_log(self, outfolder='.', basename='image'): + + if outfolder[-1] != '/': + outfolder += '/' + if not os.path.exists(outfolder): + os.makedirs(outfolder) + + self.make_waypoints = True + self.frame_num = 0 + self.outfolder = outfolder + self.basename = basename + self.mylogger('# Paicos image log file', mode='w', command=False) + self.mylogger(outfolder, command=False) + self.mylogger(basename, command=False) + + image_file = ImageWriter(self.image_creator, basedir=outfolder, + basename=f'{basename}_logfile_start') + image_file.finalize() + self.mylogger(image_file.filename, command=False) + self.logging = True + + if self.make_waypoints: + waypoint_folder = f'{self.outfolder}{self.basename}_waypoints/' + self._make_waypoint(waypoint_folder, self.basename, 0) + + def read_log(self, filename=None, outfolder='.', basename='image'): + self.make_waypoints = False + if filename is None: + filename = f'{outfolder}/{basename}.log' + with open(filename, 'r') as f: + lines = f.readlines() + self.outfolder = outfolder = lines[1].strip() + self.basename = basename = lines[2].strip() + self.logfile_start = lines[3].strip() + self.lines = lines + self.commands = self.lines[4:] + self.frame_num = 0 + self.line_index = 4 + self.num_steps = len(self.lines) - self.line_index + self.image_creator = ImageReader(self.logfile_start) + self.first_call = False + self.make_image_creator(self.image_creator.snapnum) + + waypoint_folder = f'{self.outfolder}/{self.basename}_waypoints/' + if os.path.exists(waypoint_folder): + import glob + files = glob.glob(f'{waypoint_folder}*.hdf5') + framenums = [int(file.split('_')[-2]) for file in files] + snapnums = [int(file[-8:-5]) for file in files] + index = np.argsort(framenums) + files = [files[ii] for ii in index] + self.waypoints = {} + framenums = np.array(framenums)[index] + snapnums = np.array(snapnums)[index] + for ii, frame_num in enumerate(framenums): + self.waypoints[frame_num] = files[ii] + + def resume_from_waypoint(self, frame_num): + self.frame_num = frame_num + self.line_index = self.frame_num + 4 + + self.image_creator = ImageReader(self.waypoints[frame_num]) + self.first_call = False + self.make_image_creator(self.image_creator.snapnum) + + def resume_from_image(self, filename, frame_num): + raise RuntimeError('not implemented, try resume_from_waypoint instead') + + @property + def next_action(self): + return self.lines[self.line_index].strip() + + @property + def previous_action(self): + if self.line_index == 4: + return None + else: + return self.lines[self.line_index - 1].strip() + + def step(self, verbose=False): + line = self.next_action + parts = line.split(',') + command = parts[0] + + im_creator = self.image_creator + + if verbose: + print(line) + if command == 'center_on_sub' or command == 'center_on_group': + # cx, cy, cz = float(parts[2]), float(parts[3]), float(parts[4]) + # new_center = np.array([cx, cy, cz]) * im_creator.center.uq + raise RuntimeError("not implemented") + # self.reset_center(new_center) + elif command == 'zoom': + zoom_fac = float(parts[1]) + if zoom_fac > self.zoom_max or (1 / zoom_fac) > self.zoom_max: + kk = find_zoom_sub_divisions(zoom_fac, self.zoom_max) + new_zoom_fac = zoom_fac**(1 / kk) + for _ in range(kk): + self.zoom(new_zoom_fac) + else: + self.zoom(zoom_fac) + elif command == 'width': + self.width(float(parts[1])) + elif command == 'height': + self.height(float(parts[1])) + elif command == 'depth': + self.depth(float(parts[1])) + elif command == 'rotate_around_perp_vector1': + angle = float(parts[1]) + if abs(angle) > self.max_angle: + kk = find_angle_subdivisions(angle, self.max_angle) + for _ in range(kk): + self.rotate_around_perp_vector1(angle / kk) + else: + self.rotate_around_perp_vector1(angle) + + elif command == 'rotate_around_perp_vector2': + angle = float(parts[1]) + if abs(angle) > self.max_angle: + kk = find_angle_subdivisions(angle, self.max_angle) + for _ in range(kk): + self.rotate_around_perp_vector2(angle / kk) + else: + self.rotate_around_perp_vector2(angle) + elif command == 'rotate_around_normal_vector': + angle = float(parts[1]) + if abs(angle) > self.max_angle: + kk = find_angle_subdivisions(angle, self.max_angle) + for _ in range(kk): + self.rotate_around_normal_vector(angle / kk) + else: + self.rotate_around_normal_vector(angle) + elif command == 'half_resolution': + self.half_resolution() + elif command == 'double_resolution': + self.double_resolution() + elif command == 'move_center': + diff_x, diff_y, diff_z = float( + parts[1]), float(parts[2]), float(parts[3]) + diff = np.array([diff_x, diff_y, diff_z]) * im_creator.center.uq + self.move_center(diff) + elif command == 'move_center_sim_coordinates': + diff_x, diff_y, diff_z = float( + parts[1]), float(parts[2]), float(parts[3]) + diff = np.array([diff_x, diff_y, diff_z]) * im_creator.center.uq + self.move_center_sim_coordinates(diff) + elif command == 'reset_center': + self.reset_center() + elif command == 'change_snapnum': + self.change_snapnum(int(parts[1])) + else: + raise RuntimeError(f'unknown command in log file!\n{line}\n{command}') + + if not self.logging: + self.frame_num += 1 + self.line_index += 1 + + def _make_waypoint(self, waypoint_folder, basename, frame_num): + + image_file = ImageWriter(self.image_creator, basedir=waypoint_folder, + basename=f'waypoint_{basename}_frame_{frame_num}') + image_file.finalize() + + def expand_log(self, zoom_max=1.02, max_angle=1, verbose=False, outfolder=None, basename=None): + self.zoom_max = zoom_max + self.max_angle = max_angle + + self.read_log(outfolder=self.outfolder, basename=self.basename) + + if outfolder is not None: + self.outfolder = outfolder + if basename is not None: + self.basename = basename + + self.basename = self.basename + '_expanded' + + self.create_log(outfolder=self.outfolder, basename=self.basename) + + for ii in range(4, len(self.lines)): + self.step(verbose=verbose) + + def change_snapnum(self, snapnum): + self.make_image_creator(snapnum) + + if self.logging: + self.mylogger(f'change_snapnum,{snapnum}') + + def zoom(self, factor): + im_creator = self.image_creator + im_creator.zoom(factor) + if self.logging: + self.mylogger(f'zoom,{factor}') + + def reset_center(self): + self.image_creator.reset_center() + if self.logging: + self.mylogger("reset_center") + + # def center_on_sub(self, sub_id): + # im_creator = self.image_creator + # new_center = im_creator.snap.Cat.Sub['SubhaloPos'][sub_id].T + # im_creator._diff_center += new_center - im_creator._center + # im_creator.center = new_center.copy + + # def center_on_group(self, group_id): + # im_creator = self.image_creator + # new_center = im_creator.snap.Cat.Group['GroupPos'][group_id].T + # im_creator._diff_center += new_center - im_creator._center + # im_creator.center = new_center.copy + + def width(self, arg): + im_creator = self.image_creator + im_creator.width = arg * im_creator.width.uq + if self.logging: + self.mylogger(f'width,{arg}') + + def height(self, arg): + im_creator = self.image_creator + im_creator.height = arg * im_creator.height.uq + if self.logging: + self.mylogger(f'height,{arg}') + + def depth(self, arg): + im_creator = self.image_creator + im_creator.depth = arg * im_creator.depth.uq + if self.logging: + self.mylogger(f'depth,{arg}') + + def half_resolution(self): + im_creator = self.image_creator + im_creator.half_resolution + if self.logging: + self.mylogger('half_resolution') + + def double_resolution(self): + im_creator = self.image_creator + im_creator.double_resolution + if self.logging: + self.mylogger('double_resolution') + + def rotate_around_perp_vector1(self, arg): + im_creator = self.image_creator + im_creator.orientation.rotate_around_perp_vector1(degrees=arg) + if self.logging: + self.mylogger(f'rotate_around_perp_vector1,{arg}') + + def rotate_around_perp_vector2(self, arg): + im_creator = self.image_creator + im_creator.orientation.rotate_around_perp_vector2(degrees=arg) + if self.logging: + self.mylogger(f'rotate_around_perp_vector2,{arg}') + + def rotate_around_normal_vector(self, arg): + im_creator = self.image_creator + im_creator.orientation.rotate_around_normal_vector(degrees=arg) + if self.logging: + self.mylogger(f'rotate_around_normal_vector,{arg}') + + def move_center(self, diff): + im_creator = self.image_creator + im_creator.move_center_along_perp_vector1(diff[0]) + im_creator.move_center_along_perp_vector2(diff[1]) + im_creator.move_center_along_normal_vector(diff[2]) + if self.logging: + self.mylogger(f"move_center,{diff[0].value},{diff[1].value},{diff[2].value}") + + def move_center_sim_coordinates(self, diff): + im_creator = self.image_creator + new_center = diff + im_creator.center + im_creator.center = new_center + im_creator._diff_center += diff + if self.logging: + self.mylogger(f"move_center_sim_coordinates,{diff[0].value},{diff[1].value},{diff[2].value}") diff --git a/paicos/readers/arepo_catalog.py b/paicos/readers/arepo_catalog.py index ae7a73d..7fcc33f 100755 --- a/paicos/readers/arepo_catalog.py +++ b/paicos/readers/arepo_catalog.py @@ -220,12 +220,15 @@ def load_all_data(self): if not hasattr(self.Sub[key], 'unit'): del self.Sub[key] - def save_new_catalog(self, basename, single_precision=False): + def save_new_catalog(self, basename, basedir=None, single_precision=False): """ Save a new catalog containing only the currently loaded variables. Useful for reducing datasets to smaller sizes. """ - writer = PaicosWriter(self, self.basedir, basename, 'w') + if basedir is None: + writer = PaicosWriter(self, self.basedir, basename, 'w') + else: + writer = PaicosWriter(self, basedir, basename, 'w') for key in self.Group: Ngroups_Total = self.Group[key].shape[0] diff --git a/paicos/readers/arepo_snap.py b/paicos/readers/arepo_snap.py index 52a3f32..1134656 100755 --- a/paicos/readers/arepo_snap.py +++ b/paicos/readers/arepo_snap.py @@ -845,12 +845,15 @@ def radial_select(self, center, r_max, r_min=0.0, parttype=None): return selected_snap - def save_new_snapshot(self, basename, single_precision=False): + def save_new_snapshot(self, basename, basedir=None, single_precision=False): """ Save a new snapshot containing the currently loaded (derived) variables. Useful for reducing datasets to smaller sizes. """ - writer = PaicosWriter(self, self.basedir, basename, 'w') + if basedir is None: + writer = PaicosWriter(self, self.basedir, basename, 'w') + else: + writer = PaicosWriter(self, basedir, basename, 'w') new_npart = [0] * self.nspecies for key in self: diff --git a/paicos/readers/paicos_readers.py b/paicos/readers/paicos_readers.py index ff46608..cb9912a 100644 --- a/paicos/readers/paicos_readers.py +++ b/paicos/readers/paicos_readers.py @@ -640,6 +640,11 @@ def __init__(self, basedir='.', snapnum=None, basename="projection", load_all=Tr self.direction = direction = f['image_info'].attrs['direction'] self.image_creator = f['image_info'].attrs['image_creator'] + if 'npix_width' in f['image_info'].attrs: + self.npix = self.npix_width = f['image_info'].attrs['npix_width'] + if 'npix_height' in f['image_info'].attrs: + self.npix_height = f['image_info'].attrs['npix_height'] + # Recreate orientation object if it is saved if direction == 'orientation': normal_vector = f['image_info'].attrs['normal_vector'] @@ -670,10 +675,16 @@ def __init__(self, basedir='.', snapnum=None, basename="projection", load_all=Tr self.area = self._im.area self.volume = self._im.volume - if len(list(self.keys())) > 0: - arr_shape = self[list(self.keys())[0]].shape - self.npix = self.npix_width = arr_shape[0] - self.npix_height = arr_shape[1] + if hasattr(self, 'npix_width') and hasattr(self, 'npix_height'): + pass + else: + if len(list(self.keys())) > 0: + arr_shape = self[list(self.keys())[0]].shape + self.npix = self.npix_width = arr_shape[1] + self.npix_height = arr_shape[0] + else: + import warnings + warnings.warn('Unable to read pixel dimensions') self.area_per_pixel = self.area / (self.npix_width * self.npix_height) self.volume_per_pixel = self.volume / (self.npix_width * self.npix_height) diff --git a/paicos/trees/bvh_gpu.py b/paicos/trees/bvh_gpu.py index fe03284..65d2dd6 100644 --- a/paicos/trees/bvh_gpu.py +++ b/paicos/trees/bvh_gpu.py @@ -412,6 +412,8 @@ def __init__(self, positions, sizes, threadsperblock=32): self._pos = self._pos[self.sort_index, :] self._pos_uint = self._pos_uint[self.sort_index, :] + del self._pos_uint + # Allocate arrays self.children = -1 * cp.ones((self.num_internal_nodes, 2), dtype=int) self.parents = -1 * cp.ones(self.num_leafs_and_nodes, dtype=int) @@ -419,8 +421,9 @@ def __init__(self, positions, sizes, threadsperblock=32): # This sets the parent and children properties generateHierarchy[blocks_nodes, threadsperblock](self.morton_keys, self.children, self.parents) + del self.morton_keys - # Set the boundaries for the leafs + # Set the boundaries for the leafs self.bounds = cp.zeros( (self.num_leafs_and_nodes, 3, 2), dtype=cp.uint64) @@ -464,6 +467,30 @@ def nearest_neighbor(self, query_points, threadsperblock=32): self._to_tree_coordinates(query_points), dists, ids) return dists / self.conversion_factor, self._tree_node_ids_to_data_ids(ids) + def __del__(self): + """ + Clean up like this? Apparently quite sporadic + when __del__ is called. + """ + self.release_gpu_memory() + + def release_gpu_memory(self): + if hasattr(self, '_pos'): + del self._pos + if hasattr(self, '_pos_uint'): + del self._pos_uint + if hasattr(self, 'morton_keys'): + del self.morton_keys + if hasattr(self, 'sort_index'): + del self.sort_index + if hasattr(self, 'children'): + del self.children + if hasattr(self, 'parents'): + del self.parents + if hasattr(self, 'bounds'): + del self.bounds + cp._default_memory_pool.free_all_blocks() + if __name__ == '__main__': import paicos as pa diff --git a/paicos/writers/arepo_image.py b/paicos/writers/arepo_image.py index f2a515a..19e8b5f 100644 --- a/paicos/writers/arepo_image.py +++ b/paicos/writers/arepo_image.py @@ -84,6 +84,8 @@ def __init__(self, image_creator, basedir, basename="projection", mode='w'): util.save_dataset(file, 'extent', self.extent, group='image_info') file['image_info'].attrs['direction'] = self.direction file['image_info'].attrs['image_creator'] = str(image_creator) + file['image_info'].attrs['npix_width'] = image_creator.npix_width + file['image_info'].attrs['npix_height'] = image_creator.npix_height if self.direction == 'orientation': file['image_info'].attrs['normal_vector'] = self.orientation.normal_vector diff --git a/tests/cuda-gpu/test_gpu_binary_tree.py b/tests/cuda-gpu/test_gpu_binary_tree.py index 0cf35f4..9949d88 100644 --- a/tests/cuda-gpu/test_gpu_binary_tree.py +++ b/tests/cuda-gpu/test_gpu_binary_tree.py @@ -21,10 +21,9 @@ def test_gpu_binary_tree(): index = pa.util.get_index_of_cubic_region_plus_thin_layer( snap['0_Coordinates'], center, widths, - 2.0 * snap['0_Diameters'], snap.box) + snap['0_Diameters'], snap.box) snap = snap.select(index, parttype=0) - pos_cpu = snap['0_Coordinates'] pos = cp.array(pos_cpu) sizes = cp.array(snap['0_Diameters'])