From 6c17ae8a407a05ef3ab7ddcd9e77335c5e8cc31f Mon Sep 17 00:00:00 2001
From: daikitag <48062118+daikitag@users.noreply.github.com>
Date: Wed, 23 Aug 2023 20:05:45 +0900
Subject: [PATCH] CODE: Docstring

Modify docstring in the Python codes for documentation
---
 tstrait/genetic_value.py        |  2 +-
 tstrait/simulate_environment.py |  2 +-
 tstrait/trait_model.py          | 34 +++++++++++++++++----------------
 3 files changed, 20 insertions(+), 18 deletions(-)

diff --git a/tstrait/genetic_value.py b/tstrait/genetic_value.py
index 7ce5d2f..3c6c2e6 100644
--- a/tstrait/genetic_value.py
+++ b/tstrait/genetic_value.py
@@ -278,7 +278,7 @@ def sim_genetic(ts, trait_df, alpha=0, random_seed=None):
 
     Examples
     --------
-    See :ref:`genetic` for worked examples.
+    See :ref:`genetic_value` for worked examples.
     """
 
     ts = _check_instance(ts, "ts", tskit.TreeSequence)
diff --git a/tstrait/simulate_environment.py b/tstrait/simulate_environment.py
index 2db7333..273ef7b 100644
--- a/tstrait/simulate_environment.py
+++ b/tstrait/simulate_environment.py
@@ -105,7 +105,7 @@ def sim_env(genetic_df, h2, random_seed=None):
 
     Examples
     --------
-    See :ref:`environment` for worked examples.
+    See :ref:`environment_noise` for worked examples.
     """
     genetic_df = _check_dataframe(
         genetic_df, ["trait_id", "individual_id", "genetic_value"], "genetic_df"
diff --git a/tstrait/trait_model.py b/tstrait/trait_model.py
index c172119..bb176aa 100644
--- a/tstrait/trait_model.py
+++ b/tstrait/trait_model.py
@@ -14,7 +14,7 @@ class TraitModel(metaclass=ABCMeta):
     """
     Superclass of the trait model.
 
-    Parameters
+    Attributes
     ----------
     name : str
         Name of the trait model.
@@ -479,64 +479,66 @@ def trait_model(distribution, **kwargs):
 
     Examples
     --------
+    >>> import tstrait
+
     Constructing a normal distribution trait model with mean :math:`0` and
     variance :math:`1`.
 
+    >>> import tstrait
     >>> model = tstrait.trait_model(distribution="normal", mean=0, var=1)
     >>> model.name
-    "normal"
+    'normal'
 
     Constructing a student's t-distribution trait model with mean :math:`0`,
     variance :math:`1` and degrees of freedom :math:`1`.
 
     >>> model = tstrait.trait_model(distribution="t", mean=0, var=1, df=1)
     >>> model.name
-    "t"
+    't'
 
     Constructing a fixed value trait model with value :math:`1`.
 
     >>> model = tstrait.trait_model(distribution="fixed", value=1)
     >>> model.name
-    "fixed"
+    'fixed'
 
     Constructing an exponential distribution trait model with scale
     :math:`1`.
 
     >>> model = tstrait.trait_model(distribution="exponential", scale=1)
     >>> model.name
-    "exponential"
+    'exponential'
 
     Constructing an exponential distribution trait model with scale
     :math:`1`, and enable simulation of negative values.
 
-    >>> model = tstrait.trait_model(distribution="exponential", scale=1,
-    >>> ...                         negative=True)
+    >>> model = tstrait.trait_model(distribution="exponential", scale=1, \
+                                    negative=True)
 
     Constructing a gamma distribution trait model with shape :math:`1`
     and scale :math:`2`.
 
-    >>> model = tstrait.trait_model(distribution="gamma", shape=1,
-    >>> ...                         scale=2)
+    >>> model = tstrait.trait_model(distribution="gamma", shape=1, scale=2)
     >>> model.name
-    "gamma"
+    'gamma'
 
     Constructing a gamma distribution trait model with shape :math:`1`,
     scale :math:`2`, and allow simulation of negative values.
 
-    >>> model = tstrait.trait_model(distribution="gamma", shape=1,
-    >>> ...                         scale=2, negative=True)
+    >>> model = tstrait.trait_model(distribution="gamma", shape=1, scale=2, \
+                                    negative=True)
     >>> model.name
-    "gamma"
+    'gamma'
 
     Constructing a multivariate normal distribution trait model with
     mean vector :math:`[0, 0]` and covariance matrix being an
     identity matrix.
 
     >>> import numpy as np
-    >>> model = tstrait.trait_model(distribution="multi_normal",
-    >>> ...                         mean=np.zeros(2), cov=np.eye(2))
+    >>> model = tstrait.trait_model(distribution="multi_normal", \
+                                    mean=np.zeros(2), cov=np.eye(2))
     >>> model.name
-    "multi_normal"
+    'multi_normal'
     >>> model.num_trait
     2
     """