Get Started
===========

Welcome to JittorGeometric! This guide will walk you through the basic steps to get started with graph neural networks using this library.

Quick Start
-----------

Let's start by building a simple Graph Neural Network (GNN) model using `jittor_geometric`.

Step 1: Import Libraries
------------------------

First, import the necessary libraries:

.. code-block:: python

    import jittor as jt
    from jittor import nn
    from jittor_geometric.nn import GCNConv
    from jittor_geometric.datasets import Planetoid

Step 2: Load a Dataset
----------------------

We will use the popular `Planetoid` dataset (e.g., Cora) for this example:

.. code-block:: python

    dataset = Planetoid(root='your_path', name='Cora')
    data = dataset[0]  # Getting the first graph

    # Prepare data
    from jittor_geometric.ops import cootocsr, cootocsc, gcn_norm

    edge_index, edge_weight = data.edge_index, data.edge_attr 
    edge_index, edge_weight = gcn_norm(edge_index, edge_weight, v_num, improved=False, add_self_loops=True)
    with jt.no_grad():
        data.csc = cootocsc(edge_index, edge_weight, v_num)
        data.csr = cootocsr(edge_index, edge_weight, v_num)

Step 3: Define a Simple GCN Model
---------------------------------

Now, let's define a basic Graph Convolutional Network (GCN) model:

.. code-block:: python

    class GCNModel(jt.Module):
        def __init__(self, dataset, dropout=0.8):
            super(GCNModel, self).__init__()
            self.conv1 = GCNConv(in_channels=dataset.num_features, out_channels=256, spmm=args.spmm)
            self.conv2 = GCNConv(in_channels=256, out_channels=dataset.num_classes, spmm=args.spmm)
            self.dropout = dropout

        def execute(self):
            x, csc, csr = data.x, data.csc, data.csr
            x = nn.relu(self.conv1(x, csc, csr))
            x = nn.dropout(x, self.dropout, is_train=self.training)
            x = self.conv2(x, csc, csr)
            return nn.log_softmax(x, dim=1)

Step 4: Training the Model
--------------------------

Let's train the model on the dataset:

.. code-block:: python

    # Initialize the model
    model = GCNModel(dataset)

    # Set optimizer
    optimizer = nn.Adam(params=model.parameters(), lr=0.001, weight_decay=5e-4) 

    # Training loop
    for epoch in range(200):
        model.train()
        pred = model()[data.train_mask]
        label = data.y[data.train_mask]
        loss = nn.nll_loss(pred, label)
        optimizer.step(loss)

Step 5: Evaluate the Model
--------------------------

After training, evaluate the model's performance:

.. code-block:: python

    model.eval()
    out = model()
    pred, _ = jt.argmax(out, dim=1)
    y_test = data.y[data.test_mask]
    accuracy = pred.equal(y_test).sum().item() / data.test_mask.sum().item()
    print(f'Accuracy: {accuracy.item() * 100:.2f}%')

Congratulations, you have successfully trained and tested a GNN model using `jittor_geometric`!

Next Steps
----------

- Explore more datasets: `Planetoid`, `Cora`, `Citeseer`, etc.
- Try other graph neural network layers like `SAGEConv`, `GATConv`, etc.
- Check out the documentation for more advanced features.