From ede311ec1cf5abd39a21028a91273e0225ce74b0 Mon Sep 17 00:00:00 2001
From: RemDelaporteMathurin <rdelaportemathurin@gmail.com>
Date: Tue, 4 Jun 2024 11:31:01 +0200
Subject: [PATCH] same for sympy Expression

---
 docs/source/userguide/temperature.rst | 18 +++++++++++++-----
 festim/generic_simulation.py          |  5 +++--
 test/simulation/test_initialise.py    | 20 +++++++++++++++++---
 3 files changed, 33 insertions(+), 10 deletions(-)

diff --git a/docs/source/userguide/temperature.rst b/docs/source/userguide/temperature.rst
index 98afc448f..45efbdd9d 100644
--- a/docs/source/userguide/temperature.rst
+++ b/docs/source/userguide/temperature.rst
@@ -3,7 +3,8 @@ Temperature
 ===========
 
 Definition of a temperature field or problem is essential for hydrogen transport 
-and FESTIM as a whole. 
+and FESTIM as a whole.
+Regardless of how you define the temperature of the problem, it is passed to the :code:`T` attribute of the :class:`festim.Simulation` object.
 
 ----------------------
 Analytical expressions
@@ -16,7 +17,7 @@ The temperature can be defined as a constant value in Kelvin (K):
     my_temperature = 300
 
 
-Temperature can also be defined as an expression of time and/or space using the :class:`festim.Temperature`.
+Temperature can also be defined as an expression of time and/or space.
 For example:
 
 .. math::
@@ -29,7 +30,7 @@ would be passed to FESTIM as:
 
     from festim import x, t
 
-    my_temp = Temperature(300+2*x+3*t)
+    my_temp = 300 + 2*x + 3*t
 
 More complex expressions can be expressed with sympy:
 
@@ -44,9 +45,16 @@ would be passed to FESTIM as:
     from festim import x, t
     import sympy as sp
 
-    my_temp = Temperature(sp.exp(x)*sp.sin(t))
+    my_temp = sp.exp(x) * sp.sin(t)
 
-For more details, see :class:`festim.Temperature`.
+Conditional expressions are also possible:
+
+.. code-block:: python
+
+    from festim import x, t
+    import sympy as sp
+
+    my_temp = sp.Piecewise((400, F.t < 10), (300, True))
 
 ---------------------------
 From a heat transfer solver
diff --git a/festim/generic_simulation.py b/festim/generic_simulation.py
index 354bf3805..6723efae3 100644
--- a/festim/generic_simulation.py
+++ b/festim/generic_simulation.py
@@ -2,6 +2,7 @@
 from festim.h_transport_problem import HTransportProblem
 from fenics import *
 import numpy as np
+import sympy as sp
 import warnings
 
 
@@ -25,7 +26,7 @@ class Simulation:
             None.
         settings (festim.Settings, optional): The model's settings.
             Defaults to None.
-        temperature (int, float, festim.Temperature, optional): The model's
+        temperature (int, float, sympy.Expr, festim.Temperature, optional): The model's
             temperature. Can be an expression or a heat transfer model.
             Defaults to None.
         initial_conditions (list of festim.InitialCondition, optional):
@@ -174,7 +175,7 @@ def T(self, value):
             self._T = value
         elif value is None:
             self._T = value
-        elif isinstance(value, (int, float)):
+        elif isinstance(value, (int, float, sp.Expr)):
             self._T = festim.Temperature(value)
 
     def attribute_source_terms(self):
diff --git a/test/simulation/test_initialise.py b/test/simulation/test_initialise.py
index 82521c7de..c5a3d43d3 100644
--- a/test/simulation/test_initialise.py
+++ b/test/simulation/test_initialise.py
@@ -1,6 +1,7 @@
 import festim as F
 from pathlib import Path
 import pytest
+import sympy as sp
 
 
 def test_initialise_changes_nb_of_sources():
@@ -149,11 +150,24 @@ def test_cartesian_and_surface_flux_warning(quantity, sys):
         my_model.initialise()
 
 
-@pytest.mark.parametrize("value", [100, 0, 100.0])
-def test_initialise_temp_as_number(value):
+@pytest.mark.parametrize(
+    "value",
+    [
+        100,
+        0,
+        100.0,
+        0.0,
+        100 + 1 * F.x,
+        0 + 1 * F.x,
+        F.x,
+        F.t,
+        sp.Piecewise((400, F.t < 10), (300, True)),
+    ],
+)
+def test_initialise_temp_as_number_or_sympy(value):
     """
     Creates a Simulation object and checks that the T attribute
-    can be given as an int or float
+    can be given as an int, a float or a sympy Expr
     """
     # build
     my_model = F.Simulation()