From 6e44cb85c4f204cebc83143bfb3670b4d4d6773a Mon Sep 17 00:00:00 2001
From: Ardavan Oskooi <ardavano@google.com>
Date: Thu, 15 Dec 2022 14:55:52 -0800
Subject: [PATCH] Add a function to return the number of elapsed timesteps in
 `Simulation` class object (#2337)

* add a function to return the number of elapsed timesteps in Simulation

* specify runtime based on number of timesteps and sim.fields.dt

* Update doc/docs/Python_User_Interface.md.in

Co-authored-by: Steven G. Johnson <stevenj@mit.edu>
---
 doc/docs/Python_User_Interface.md.in |  1 +
 python/simulation.py                 |  7 +++++++
 python/tests/test_simulation.py      | 11 +++++++++++
 3 files changed, 19 insertions(+)

diff --git a/doc/docs/Python_User_Interface.md.in b/doc/docs/Python_User_Interface.md.in
index 00b6a0808..53c6f096d 100644
--- a/doc/docs/Python_User_Interface.md.in
+++ b/doc/docs/Python_User_Interface.md.in
@@ -85,6 +85,7 @@ or set the output folder, with these methods of the `Simulation` class:
 The `Simulation` class provides the following time-related methods:
 
 @@ Simulation.meep_time @@
+@@ Simulation.timestep @@
 @@ Simulation.print_times @@
 @@ Simulation.time_spent_on @@
 @@ Simulation.mean_time_spent_on @@
diff --git a/python/simulation.py b/python/simulation.py
index af0f2407c..cbd3d2e71 100644
--- a/python/simulation.py
+++ b/python/simulation.py
@@ -2429,6 +2429,13 @@ def meep_time(self):
             self.init_sim()
         return self.fields.time()
 
+    def timestep(self) -> int:
+        """Return the number of elapsed timesteps."""
+
+        if self.fields is None:
+            self.init_sim()
+        return self.fields.t
+
     def round_time(self):
         if self.fields is None:
             self.init_sim()
diff --git a/python/tests/test_simulation.py b/python/tests/test_simulation.py
index a0c3b3b7a..271dc1617 100644
--- a/python/tests/test_simulation.py
+++ b/python/tests/test_simulation.py
@@ -828,6 +828,17 @@ def test_iterable_as_v3(self):
             self.assertAlmostEqual(pt1, expected)
             self.assertAlmostEqual(pt2, expected)
 
+    def test_time(self):
+        sim = self.init_simple_simulation()
+
+        num_timesteps = 212
+        sim.init_sim()
+        end_t = sim.fields.dt * num_timesteps
+        sim.run(until=end_t)
+
+        self.assertAlmostEqual(sim.meep_time(), end_t)
+        self.assertEqual(sim.timestep(), num_timesteps)
+
 
 if __name__ == "__main__":
     unittest.main()