Пример #1
0
    def _generate_kdtree(self, fname):
        from yt.utilities.lib.cykdtree import PyKDTree

        if fname is not None:
            if os.path.exists(fname):
                mylog.info("Loading KDTree from %s", os.path.basename(fname))
                kdtree = PyKDTree.from_file(fname)
                if kdtree.data_version != self.ds._file_hash:
                    mylog.info("Detected hash mismatch, regenerating KDTree")
                else:
                    self._kdtree = kdtree
                    return
        positions = []
        for data_file in self.data_files:
            for _, ppos in self.io._yield_coordinates(
                    data_file, needed_ptype=self.ds._sph_ptypes[0]):
                positions.append(ppos)
        if positions == []:
            self._kdtree = None
            return
        positions = np.concatenate(positions)
        mylog.info("Allocating KDTree for %s particles", positions.shape[0])
        num_neighbors = getattr(self.ds, "num_neighbors", 32)
        self._kdtree = PyKDTree(
            positions.astype("float64"),
            left_edge=self.ds.domain_left_edge,
            right_edge=self.ds.domain_right_edge,
            periodic=np.array(self.ds.periodicity),
            leafsize=2 * int(num_neighbors),
            data_version=self.ds._file_hash,
        )
        if fname is not None:
            self._kdtree.save(fname)
Пример #2
0
class SPHParticleIndex(ParticleIndex):
    def _initialize_index(self):
        ds = self.dataset

        ds._file_hash = self._generate_hash()

        if hasattr(self.io, '_generate_smoothing_length'):
            self.io._generate_smoothing_length(self.data_files, self.kdtree)

        super(SPHParticleIndex, self)._initialize_index()

    def _generate_kdtree(self, fname):
        from yt.utilities.lib.cykdtree import PyKDTree
        if fname is not None:
            if os.path.exists(fname):
                mylog.info('Loading KDTree from %s' % os.path.basename(fname))
                kdtree = PyKDTree.from_file(fname)
                if kdtree.data_version != self.ds._file_hash:
                    mylog.info('Detected hash mismatch, regenerating KDTree')
                else:
                    self._kdtree = kdtree
                    return
        positions = []
        for data_file in self.data_files:
            for _, ppos in self.io._yield_coordinates(
                    data_file, needed_ptype=self.ds._sph_ptypes[0]):
                positions.append(ppos)
        if positions == []:
            self._kdtree = None
            return
        positions = np.concatenate(positions)
        mylog.info('Allocating KDTree for %s particles' % positions.shape[0])
        self._kdtree = PyKDTree(positions.astype('float64'),
                                left_edge=self.ds.domain_left_edge,
                                right_edge=self.ds.domain_right_edge,
                                periodic=np.array(self.ds.periodicity),
                                leafsize=2 * int(self.ds.num_neighbors),
                                data_version=self.ds._file_hash)
        if fname is not None:
            self._kdtree.save(fname)

    @property
    def kdtree(self):
        if hasattr(self, '_kdtree'):
            return self._kdtree

        ds = self.ds

        if getattr(ds, 'kdtree_filename', None) is None:
            if os.path.exists(ds.parameter_filename):
                fname = ds.parameter_filename + ".kdtree"
            else:
                # we don't want to write to disk for in-memory data
                fname = None
        else:
            fname = ds.kdtree_filename

        self._generate_kdtree(fname)

        return self._kdtree
Пример #3
0
    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