def _generate_smoothing_length(self, index): data_files = index.data_files if not self.ds.gen_hsmls: return hsml_fn = data_files[0].filename.replace(".hdf5", ".hsml.hdf5") if os.path.exists(hsml_fn): with h5py.File(hsml_fn, mode="r") as f: file_hash = f.attrs["q"] if file_hash != self.ds._file_hash: mylog.warning("Replacing hsml files.") for data_file in data_files: hfn = data_file.filename.replace(".hdf5", ".hsml.hdf5") os.remove(hfn) else: return positions = [] counts = defaultdict(int) for data_file in data_files: for _, ppos in self._yield_coordinates( data_file, needed_ptype=self.ds._sph_ptypes[0] ): counts[data_file.filename] += ppos.shape[0] positions.append(ppos) if not positions: return offsets = {} offset = 0 for fn, count in counts.items(): offsets[fn] = offset offset += count kdtree = index.kdtree positions = uconcatenate(positions)[kdtree.idx] hsml = generate_smoothing_length( positions.astype("float64"), kdtree, self.ds._num_neighbors ) dtype = positions.dtype hsml = hsml[np.argsort(kdtree.idx)].astype(dtype) mylog.warning("Writing smoothing lengths to hsml files.") for i, data_file in enumerate(data_files): si, ei = data_file.start, data_file.end fn = data_file.filename hsml_fn = data_file.filename.replace(".hdf5", ".hsml.hdf5") with h5py.File(hsml_fn, mode="a") as f: if i == 0: f.attrs["q"] = self.ds._file_hash g = f.require_group(self.ds._sph_ptypes[0]) d = g.require_dataset( "SmoothingLength", dtype=dtype, shape=(counts[fn],) ) begin = si + offsets[fn] end = min(ei, d.size) + offsets[fn] d[si:ei] = hsml[begin:end]
def _generate_smoothing_length(self, data_files, kdtree): if os.path.exists(self.hsml_filename): with open(self.hsml_filename, "rb") as f: file_hash = struct.unpack("q", f.read(struct.calcsize("q")))[0] if file_hash != self.ds._file_hash: os.remove(self.hsml_filename) else: return positions = [] for data_file in data_files: for _, ppos in self._yield_coordinates( data_file, needed_ptype=self.ds._sph_ptypes[0] ): positions.append(ppos) if positions == []: return positions = np.concatenate(positions)[kdtree.idx] hsml = generate_smoothing_length(positions, kdtree, self.ds._num_neighbors) hsml = hsml[np.argsort(kdtree.idx)] dtype = self._pdtypes["Gas"]["Coordinates"][0] with open(self.hsml_filename, "wb") as f: f.write(struct.pack("q", self.ds._file_hash)) f.write(hsml.astype(dtype).tostring())
def add_sph_fields(self, n_neighbors=32, kernel="cubic", sph_ptype="io"): """Add SPH fields for the specified particle type. For a particle type with "particle_position" and "particle_mass" already defined, this method adds the "smoothing_length" and "density" fields. "smoothing_length" is computed as the distance to the nth nearest neighbor. "density" is computed as the SPH (gather) smoothed mass. The SPH fields are added only if they don't already exist. Parameters ---------- n_neighbors : int The number of neighbors to use in smoothing length computation. kernel : str The kernel function to use in density estimation. sph_ptype : str The SPH particle type. Each dataset has one sph_ptype only. This method will overwrite existing sph_ptype of the dataset. """ mylog.info("Generating SPH fields") # Unify units l_unit = "code_length" m_unit = "code_mass" d_unit = "code_mass / code_length**3" # Read basic fields ad = self.all_data() pos = ad[sph_ptype, "particle_position"].to(l_unit).d mass = ad[sph_ptype, "particle_mass"].to(m_unit).d # Construct k-d tree kdtree = PyKDTree( pos.astype("float64"), left_edge=self.domain_left_edge.to_value(l_unit), right_edge=self.domain_right_edge.to_value(l_unit), periodic=self.periodicity, leafsize=2 * int(n_neighbors), ) order = np.argsort(kdtree.idx) def exists(fname): if (sph_ptype, fname) in self.derived_field_list: mylog.info("Field ('%s','%s') already exists. Skipping", sph_ptype, fname) return True else: mylog.info("Generating field ('%s','%s')", sph_ptype, fname) return False data = {} # Add smoothing length field fname = "smoothing_length" if not exists(fname): hsml = generate_smoothing_length(pos[kdtree.idx], kdtree, n_neighbors) hsml = hsml[order] data[(sph_ptype, "smoothing_length")] = (hsml, l_unit) else: hsml = ad[sph_ptype, fname].to(l_unit).d # Add density field fname = "density" if not exists(fname): dens = estimate_density( pos[kdtree.idx], mass[kdtree.idx], hsml[kdtree.idx], kdtree, kernel_name=kernel, ) dens = dens[order] data[(sph_ptype, "density")] = (dens, d_unit) # Add fields self._sph_ptypes = (sph_ptype, ) self.index.update_data(data) self.num_neighbors = n_neighbors