How to use the schnetpack.Properties.R function in schnetpack

def test_main_file_parser(main_path, targets_main):
    main_parser = OrcaMainFileParser(properties=OrcaMainFileParser.properties)
    main_parser.parse_file(main_path)

    results = main_parser.get_parsed()
    results[Properties.Z] = results["atoms"][0]
    results[Properties.R] = results["atoms"][1]
    results.pop("atoms", None)

    for p in targets_main:
        assert p in results

        if p == Properties.Z:
            assert np.array_equal(results[p], targets_main[p])
        else:
            assert np.allclose(results[p], targets_main[p])

def _generate_input(self, system):
        """
        Function to extracts neighbor lists, atom_types, positions e.t.c. from the system and generate a properly
        formatted input for the schnetpack model.

        Args:
            system (schnetpack.md.System): System object containing current state of the simulation.

        Returns:
            dict(torch.Tensor): Schnetpack inputs in dictionary format.
        """
        positions, atom_types, atom_masks = self._get_system_molecules(system)
        neighbors, neighbor_mask = self._get_system_neighbors(system)

        inputs = {
            Properties.R: positions,
            Properties.Z: atom_types,
            Properties.atom_mask: atom_masks,
            Properties.cell: None,
            Properties.cell_offset: None,
            Properties.neighbors: neighbors,
            Properties.neighbor_mask: neighbor_mask,
        }

        return inputs

end = i + len(d['valid'])
            idcs = np.where(d['valid'])[0]
            if len(idcs) < 1:
                i = end
                continue
            # shrink stats
            idx_id = stat_heads.index('id')
            idx_known = stat_heads.index('known')
            new_stats = stats[:, start:end]
            if 'new' in args.store and args.model_path is not None:
                idcs = idcs[np.where(new_stats[idx_known, idcs] == 0)[0]]
            new_stats = new_stats[:, idcs]
            new_stats[idx_id] = np.arange(len(new_stats[idx_id]))  # adjust ids
            shrunk_stats = np.hstack((shrunk_stats, new_stats))
            # shrink positions and atomic numbers
            shrunk_res[key] = {Properties.R: d[Properties.R][idcs],
                               Properties.Z: d[Properties.Z][idcs]}
            # store connectivity matrices if desired
            if 'connectivity' in args.store:
                shrunk_res[key].update(
                    {'connectivity': [d['connectivity'][k] for k in idcs]})
            i = end

        shrunk_res['stats'] = shrunk_stats
        res = shrunk_res

    # store results in new database
    # get filename that is not yet taken for db
    if os.path.isfile(target_db):
        file_name, _ = os.path.splitext(target_db)
        expand = 0
        while True:

def get_positions(self, mol_idx=0, replica_idx=None):
        """
        Auxiliary routine for getting the positions of specific molecules and replicas.

        Args:
            mol_idx (int): Index of the molecule to extract, by default uses the first molecule (mol_idx=0)
            replica_idx (int): Replica of the molecule to extract (e.g. for ring polymer molecular dynamics). If
                               replica_idx is set to None (default), the centroid is returned if multiple replicas are
                               present.

        Returns:
            np.array: N_steps x N_atoms x 3 array containing the atom positions of the simulation in atomic units.
        """
        return self.get_property(
            Properties.R, mol_idx=mol_idx, replica_idx=replica_idx, atomistic=True
        )

present for the Loader.
        """
        # This is for molecule streams
        structures = self.database["molecules"]

        # General database info
        # TODO: Could be moved to global attrs if available
        self.n_replicas = structures.attrs["n_replicas"]
        self.n_molecules = structures.attrs["n_molecules"]
        self.n_atoms = structures.attrs["n_atoms"]
        self.entries = structures.attrs["entries"]
        self.time_step = structures.attrs["time_step"]

        # Write to main property dictionary
        self.properties[Properties.Z] = structures.attrs["atom_types"][0, ...]
        self.properties[Properties.R] = structures[
            self.skip_initial : self.entries, ..., :3
        ]
        self.properties["velocities"] = structures[
            self.skip_initial : self.entries, ..., 3:
        ]

h, m, s = int(h), int(m), int(s)
        print(f'Time consumed: {h:d}:{m:02d}:{s:02d}')

    # sort keys in resulting dictionary
    generated = dict(sorted(generated.items()))

    # show generated molecules and print some statistics if desired
    if args.show_gen:
        ats = []
        n_total_atoms = 0
        n_molecules = 0
        for key in generated:
            n = 0
            for i in range(len(generated[key][Properties.Z])):
                at = Atoms(generated[key][Properties.Z][i],
                           positions=generated[key][Properties.R][i])
                ats += [at]
                n += 1
                n_molecules += 1
            n_total_atoms += n * key
        asv.view(ats)
        print(f'Total number of atoms placed: {n_total_atoms} '
              f'(avg {n_total_atoms / n_molecules:.2f})', flush=True)

    return generated

def get_positions(self, mol_idx=0, replica_idx=None):
        """
        Auxiliary routine for getting the positions of specific molecules and replicas.

        Args:
            mol_idx (int): Index of the molecule to extract, by default uses the first molecule (mol_idx=0)
            replica_idx (int): Replica of the molecule to extract (e.g. for ring polymer molecular dynamics). If
                               replica_idx is set to None (default), the centroid is returned if multiple replicas are
                               present.

        Returns:
            np.array: N_steps x N_atoms x 3 array containing the atom positions of the simulation in atomic units.
        """
        return self.get_property(
            Properties.R, mol_idx=mol_idx, replica_idx=replica_idx, atomistic=True
        )

def forward(self, inputs):
        """
        Args:
            inputs (dict of torch.Tensor): SchNetPack format dictionary of input tensors.

        Returns:
            torch.Tensor: Nbatch x Natoms x Nsymmetry_functions Tensor containing ACSFs or wACSFs.

        """
        positions = inputs[Properties.R]
        Z = inputs[Properties.Z]
        neighbors = inputs[Properties.neighbors]
        neighbor_mask = inputs[Properties.neighbor_mask]

        cell = inputs[Properties.cell]
        cell_offset = inputs[Properties.cell_offset]

        # Compute radial functions
        if self.RDF is not None:
            # Get atom type embeddings
            Z_rad = self.radial_Z(Z)
            # Get atom types of neighbors
            Z_ij = snn.neighbor_elements(Z_rad, neighbors)
            # Compute distances
            distances = snn.atom_distances(
                positions,

def forward(self, inputs):
        """
        Predicts the electronic spatial extent.
        """
        positions = inputs[Properties.R]
        atom_mask = inputs[Properties.atom_mask][:, :, None]

        # run prediction
        charges = self.out_net(inputs) * atom_mask
        yi = torch.norm(positions, 2, 2, keepdim=True) ** 2 * charges
        y = self.atom_pool(yi)

        # collect results
        result = {self.property: y}

        if self.contributions:
            result[self.contributions] = charges

        return result