jittor_geometric.data

Author: lusz Date: 2025-01-14 00:47:42 Description:

class jittor_geometric.data.Batch(x=None, edge_index=None, edge_attr=None, y=None, pos=None, normal=None, face=None, column_indices=None, row_offset=None, csr_edge_weight=None, row_indices=None, column_offset=None, csc_edge_weight=None, **kwargs)[source]

A data object describing a batch of graphs as one big (disconnected) graph. Inherits from jittor_geometric.data.Data.

property batch_size: int

Alias for num_graphs.

classmethod from_data_list(data_list, follow_batch=None, exclude_keys=None)[source]

Constructs a Batch object from a list of Data objects. The assignment vector batch is created on the fly. In addition, creates assignment vectors for each key in follow_batch. Will exclude any keys given in exclude_keys.

Return type:

Batch

get_example(idx)[source]

Gets the Data object at index idx. The Batch object must have been created via from_data_list() in order to be able to reconstruct the initial object.

Return type:

Data

property num_graphs: int

Returns the number of graphs in the batch.

to_data_list()[source]

Reconstructs the list of Data objects from the Batch object.

Return type:

List[Data]

class jittor_geometric.data.CSC(row_indices=None, column_offset=None, edge_weight=None)[source]
class jittor_geometric.data.CSR(column_indices=None, row_offset=None, edge_weight=None)[source]
class jittor_geometric.data.ConformerGen(**params)[source]

This class designed to generate conformers for molecules represented as SMILES strings using provided parameters and configurations. The transform method uses multiprocessing to speed up the conformer generation process.

single_process(smiles)[source]

Processes a single SMILES string to generate conformers using the specified method.

Parameters:

smiles – (str) The SMILES string representing the molecule.

Returns:

A unimolecular data representation (dictionary) of the molecule.

Raises:

ValueError – If the conformer generation method is unrecognized.

transform(smiles_list)[source]
transform_raw(atoms_list, coordinates_list)[source]
class jittor_geometric.data.Data(x=None, edge_index=None, edge_attr=None, y=None, pos=None, normal=None, face=None, column_indices=None, row_offset=None, csr_edge_weight=None, row_indices=None, column_offset=None, csc_edge_weight=None, **kwargs)[source]

A plain old python object modeling a single graph with various (optional) attributes:

Parameters:
  • x (Var, optional) – Node feature matrix with shape [num_nodes, num_node_features]. (default: None)

  • edge_index (Var.int32, optional) – Graph connectivity in COO format with shape [2, num_edges]. (default: None)

  • edge_attr (Var, optional) – Edge feature matrix with shape [num_edges, num_edge_features]. (default: None)

  • y (Var, optional) – Graph or node targets with arbitrary shape. (default: None)

  • pos (Var, optional) – Node position matrix with shape [num_nodes, num_dimensions]. (default: None)

  • normal (Var, optional) – Normal vector matrix with shape [num_nodes, num_dimensions]. (default: None)

  • face (Var.int32, optional) – Face adjacency matrix with shape [3, num_faces]. (default: None)

The data object is not restricted to these attributes and can be extented by any other additional data.

Example:

data = Data(x=x, edge_index=edge_index)
data.train_idx = jt.array([...], dtype=Var.int32)
data.test_mask = jt.array([...], dtype=Var.bool)
apply(func, *keys)[source]

Applies the function func to all Var attributes *keys. If *keys is not given, func is applied to all present attributes.

clone()[source]

Performs a deep-copy of the data object.

contains_isolated_nodes()[source]

Returns True, if the graph contains isolated nodes.

contains_self_loops()[source]

Returns True, if the graph contains self-loops.

contiguous(*keys)[source]

Ensures a contiguous memory layout for all attributes *keys. If *keys is not given, all present attributes are ensured to have a contiguous memory layout.

cpu(*keys)[source]

Copies all attributes *keys to CPU memory. If *keys is not given, the conversion is applied to all present attributes.

cuda(device=None, non_blocking=False, *keys)[source]

Copies all attributes *keys to CUDA memory. If *keys is not given, the conversion is applied to all present attributes.

debug()[source]
classmethod from_dict(dictionary)[source]

Creates a data object from a python dictionary.

property keys

Returns all names of graph attributes.

property num_edge_features

Returns the number of features per edge in the graph.

property num_edges

Returns the number of edges in the graph. For undirected graphs, this will return the number of bi-directional edges, which is double the amount of unique edges.

property num_faces

Returns the number of faces in the mesh.

property num_features

Alias for num_node_features.

property num_node_features

Returns the number of features per node in the graph.

property num_nodes

Returns or sets the number of nodes in the graph.

Note

The number of nodes in your data object is typically automatically inferred, e.g., when node features x are present. In some cases however, a graph may only be given by its edge indices edge_index. Jittor Geometric then guesses the number of nodes according to edge_index.max().item() + 1, but in case there exists isolated nodes, this number has not to be correct and can therefore result in unexpected batch-wise behavior. Thus, we recommend to set the number of nodes in your data object explicitly via data.num_nodes = .... You will be given a warning that requests you to do so.

pin_memory(*keys)[source]

Copies all attributes *keys to pinned memory. If *keys is not given, the conversion is applied to all present attributes.

to_dict()[source]
to_namedtuple()[source]
class jittor_geometric.data.Dataset(root=None, transform=None, pre_transform=None, pre_filter=None)[source]
Parameters:
  • root (string, optional) – Root directory where the dataset should be saved. (optional: None)

  • transform (callable, optional) – A function/transform that takes in an torch_geometric.data.Data object and returns a transformed version. The data object will be transformed before every access. (default: None)

  • pre_transform (callable, optional) – A function/transform that takes in an jittor_geometric.data.Data object and returns a transformed version. The data object will be transformed before being saved to disk. (default: None)

  • pre_filter (callable, optional) – A function that takes in an jittor_geometric.data.Data object and returns a boolean value, indicating whether the data object should be included in the final dataset. (default: None)

download()[source]

Downloads the dataset to the self.raw_dir folder.

get(idx)[source]

Gets the data object at index idx.

index_select(idx)[source]
indices()[source]
len()[source]
property num_edge_features

Returns the number of features per edge in the dataset.

property num_features

Alias for num_node_features.

property num_node_features

Returns the number of features per node in the dataset.

process()[source]

Processes the dataset to the self.processed_dir folder.

property processed_dir
property processed_file_names

The name of the files to find in the self.processed_dir folder in order to skip the processing.

property processed_paths

The filepaths to find in the self.processed_dir folder in order to skip the processing.

property raw_dir
property raw_file_names

The name of the files to find in the self.raw_dir folder in order to skip the download.

property raw_paths

The filepaths to find in order to skip the download.

shuffle(return_perm=False)[source]

Randomly shuffles the examples in the dataset.

Parameters:

return_perm (bool, optional) – If set to True, will additionally return the random permutation used to shuffle the dataset. (default: False)

class jittor_geometric.data.Dictionary(*, bos='[CLS]', pad='[PAD]', eos='[SEP]', unk='[UNK]', extra_special_symbols=None)[source]

A mapping from symbols to consecutive integers

add_from_file(f)[source]

Loads pre-defined dictionary symbols. If f == “default”, it will load the default atom dictionary. Otherwise, loads from a text file and adds its symbols to this instance.

add_symbol(word, n=1, overwrite=False, is_special=False)[source]

Adds a word to the dictionary

bos()[source]

Helper to get index of beginning-of-sentence symbol

eos()[source]

Helper to get index of end-of-sentence symbol

index(sym)[source]

Returns the index of the specified symbol

classmethod load(f)[source]

Loads the dictionary from a text file with the format:

` <symbol0> <count0> <symbol1> <count1> ... `

pad()[source]

Helper to get index of pad symbol

special_index()[source]
unk()[source]

Helper to get index of unk symbol

vec_index(a)[source]
class jittor_geometric.data.GraphChunk(chunks, chunk_id, v_num, global_v_num, local_mask=None, local_feature=None, local_label=None)[source]
static load(file_path)[source]

Load a GraphChunk instance from a binary file. :type file_path: str :param file_path: Path to the file from which the instance will be loaded. :return: Loaded GraphChunk instance.

save(file_path)[source]

Save the GraphChunk instance as a binary file. :type file_path: str :param file_path: Path to the file where the instance will be saved.

set_csr(column_indices, row_offset, edge_weight=None)[source]

Set the CSR (Compressed Sparse Row) representation of the graph. :type column_indices: :param column_indices: Column indices of the non-zero elements. :type row_offset: :param row_offset: Row offsets for the CSR format. :type edge_weight: :param edge_weight: Optional edge weights.

class jittor_geometric.data.InMemoryDataset(root=None, transform=None, pre_transform=None, pre_filter=None)[source]

Dataset base class for creating graph datasets which fit completely into CPU memory.

Parameters:
  • root (string, optional) – Root directory where the dataset should be saved. (default: None)

  • transform (callable, optional) – A function/transform that takes in an jittor_geometric.data.Data object and returns a transformed version. The data object will be transformed before every access. (default: None)

  • pre_transform (callable, optional) – A function/transform that takes in an jittor_geometric.data.Data object and returns a transformed version. The data object will be transformed before being saved to disk. (default: None)

  • pre_filter (callable, optional) – A function that takes in an jittor_geometric.data.Data object and returns a boolean value, indicating whether the data object should be included in the final dataset. (default: None)

static collate(data_list)[source]

Collates a python list of data objects to the internal storage format of torch_geometric.data.InMemoryDataset.

copy(idx=None)[source]
download()[source]

Downloads the dataset to the self.raw_dir folder.

get(idx)[source]

Gets the data object at index idx.

len()[source]
property num_classes

The number of classes in the dataset.

process()[source]

Processes the dataset to the self.processed_dir folder.

property processed_file_names

The name of the files to find in the self.processed_dir folder in order to skip the processing.

property raw_file_names

The name of the files to find in the self.raw_dir folder in order to skip the download.

class jittor_geometric.data.TemporalData(src=None, dst=None, t=None, msg=None, y=None, edge_ids=None, **kwargs)[source]
apply(func, *keys)[source]

Applies the function func to all Var attributes *keys. If *keys is not given, func is applied to all present attributes.

property keys
property num_edges
property num_events
property num_nodes
seq_batches(batch_size)[source]
to(device, *keys, **kwargs)[source]
train_val_test_split(val_ratio=0.15, test_ratio=0.15)[source]
train_val_test_split_w_mask()[source]
jittor_geometric.data.decide_download(url)[source]
jittor_geometric.data.download_url(url, folder, log=True)[source]

Downloads the content of an URL to a specific folder.

Parameters:
  • url (string) – The url.

  • folder (string) – The folder.

  • log (bool, optional) – If False, will not print anything to the console. (default: True)

jittor_geometric.data.extract_zip(path, folder, log=True)[source]

Extracts a zip archive to a specific folder. :type path: :param path: The path to the tar archive. :type path: string :type folder: :param folder: The folder. :type folder: string :type log: :param log: If False, will not print anything to the

console. (default: True)

```