From d3f56cf2b30277bd6a527a745a7fe444e72e8a5d Mon Sep 17 00:00:00 2001
From: obliviateandsurrender <utkarshazad98@gmail.com>
Date: Sat, 12 Aug 2023 10:51:12 -0400
Subject: [PATCH] explain `arguments`

---
 README.md                     | 14 ++++++++++++++
 quantum_datasets/generator.py | 28 ++++++++++++++++++++++++----
 2 files changed, 38 insertions(+), 4 deletions(-)

diff --git a/README.md b/README.md
index c609be8..4f798bd 100644
--- a/README.md
+++ b/README.md
@@ -20,6 +20,13 @@ import quantum_datasets as qd
 qd.qchem_data_generate("xyzfiles/q.1-1/schm-1/h2.xyz")
 ```
 
+The `qchem_data_generate` method accepts the following arguments:
+1. ``xyz_path (str)``: absolute path to the `xyz` file containing basic information regarding the molecule
+2. ``basis (str)``: basis set (STO-3G, 6-31G or CCVPDZ) used for building the molecule and the electronic hamiltonian
+3. ``bondlenths (list)``: list of additional bondlengths used for data generation using the geometry in `geom_struct`
+4. ``use_bond_struct (bool)``: whether to use bondlengths of interest in `bond_struct` for PES generation
+5. ``folder_path (str)``: absolute path for storing the generated data files
+
 For the quantum spin systems, we have to provide the name of the spin systems (for example: "Ising", "Heisenberg", "FermiHubbard" or "BoseHubbard"), and the periodicity and layout of the lattice.
 
 ``` python
@@ -27,5 +34,12 @@ import quantum_datasets as qd
 qd.qspin_data_generate("Ising", periodicity=True, layout=(1, 4), num_systems=1000)
 ```
 
+The `qspin_data_generate` method accepts the following arguments:
+1. ``sysname (str)``: type of spin model ("Ising", "Heisenberg", "FermiHubbard", "BoseHubbard") for which data has to be generated
+2. ``periodicity (bool)``: whether the lattice for the spin model is closed or open
+3. ``layout (list/tuple)``: layout for the lattice of the spin model. For example - (1, 8) or (2, 2) 
+4. ``num_systems (int)``: number of spins systems generated by varying the model parameters
+5. ``folder_path (str)``: absolute path for storing the generated data files
+
 ### Datasets
 All the datasets will be generated and written under the `datasets` folder with distinct subfolders for `qchem` and `qspin`.
\ No newline at end of file
diff --git a/quantum_datasets/generator.py b/quantum_datasets/generator.py
index 7070543..ff7670d 100644
--- a/quantum_datasets/generator.py
+++ b/quantum_datasets/generator.py
@@ -27,11 +27,19 @@
 }
 
 
-# pylint: disable=dangerous-default-value
+# pylint: disable=dangerous-default-value, line-too-long
 def qchem_data_generate(
     xyz_path, basis="STO-3G", bondlengths=[], use_bond_struct=False, folder_path=None
 ):
-    r"""Generates data for Molecular Systems"""
+    r"""Generates data for Molecular Systems
+    
+    Args:
+        xyz_path (str): absolute path to the `xyz` file containing basic information regarding the molecule
+        basis (str): basis set (STO-3G, 6-31G or CCVPDZ) used for building the molecule and the electronic hamiltonian
+        bondlenths (list): list of additional bondlengths used for data generation using the geometry in `geom_struct`
+        use_bond_struct (bool): whether to use bondlengths of interest in `bond_struct` for PES generation
+        folder_path (str): absolute path for storing the generated data files
+    """
     data_pipeline = ChemDataPipeline()
     _, mol_name, symbols, charge, geometry, gs_bond = read_xyz(xyz_path)
     geometries, gs_bonds = [geometry], [gs_bond]
@@ -59,7 +67,15 @@ def qchem_data_generate(
         gc.collect()
 
 def qspin_data_generate(sysname, periodicity, layout, num_systems=100, folder_path=None):
-    r"""Generates data for Spin Systems"""
+    r"""Generates data for Spin Systems
+
+    Args:
+        sysname (str): type of spin model (Ising, Heisenberg, FermiHubbard, BoseHubbard) for which data has to be generated
+        periodicity (bool): whether the lattice for the spin model is closed or open
+        layout (list): layout for the lattice of the spin model. For example - (1, 8) or (2, 2) 
+        num_systems (int): number of spins systems generated by varying the model parameters
+        folder_path (str): absolute path for storing the generated data files
+    """
     data_pipeline = SpinDataPipeline()
     lat = nx.grid_2d_graph(layout[0], layout[1], periodic=periodicity)
     try:
@@ -73,7 +89,11 @@ def qspin_data_generate(sysname, periodicity, layout, num_systems=100, folder_pa
     lattice = "rectangular" if layout[0] > 1 else "chain"
     sites = f"{layout[0]}x{layout[1]}"
 
-    file_name = f"datasets/qspin/{sysname.lower()}/{periodic}/{lattice}/{sites}/{sysname.lower()}_{periodic}_{lattice}_{sites}"
+    if folder_path is not None:
+        file_name = f"datasets/qspin/{sysname.lower()}/{periodic}/{lattice}/{sites}/{sysname.lower()}_{periodic}_{lattice}_{sites}"
+    else:
+        file_name = f"{folder_path}/{sysname.lower()}_{periodic}_{lattice}_{sites}"
+    
     system = spinsys(num_systems, layout, lat, periodicity)
     for _ in (pbar := tqdm(range(1))):
         data_pipeline.pipeline(sysname.lower(), system, filepath=file_name, prog_bar=pbar)