From 74060657a2c4a2c21553c19eb157676d2bad607b Mon Sep 17 00:00:00 2001
From: Martin Lang <martin.lang@mpsd.mpg.de>
Date: Sun, 14 Jul 2024 15:38:44 +0200
Subject: [PATCH] Support dict and Field for ZhangLi.u

---
 mumax3c/scripts/driver.py | 24 ++++++++++++++++++------
 mumax3c/tests/conftest.py |  5 ++---
 2 files changed, 20 insertions(+), 9 deletions(-)

diff --git a/mumax3c/scripts/driver.py b/mumax3c/scripts/driver.py
index 9635dca..ceba844 100644
--- a/mumax3c/scripts/driver.py
+++ b/mumax3c/scripts/driver.py
@@ -54,18 +54,30 @@ def driver_script(driver, system, compute=None, ovf_format="bin4", **kwargs):
 
         if system.dynamics.get(type=mm.ZhangLi):
             (zh_li_term,) = system.dynamics.get(type=mm.ZhangLi)
-            if isinstance(zh_li_term.u, df.Field):
-                raise RuntimeError("Setting Zhang Li u with a 'Fiel' is not supported.")
+            if isinstance(zh_li_term.u, df.Field) and zh_li_term.u.nvdim == 3:
+                u = zh_li_term.u
+            elif isinstance(zh_li_term.u, df.Field) and zh_li_term.u.nvdim == 1:
+                zero_field = df.Field(
+                    mesh=system.m.mesh,
+                    nvdim=1,
+                    value=0.0,
+                )
+                u = zh_li_term.u << zero_field << zero_field
             elif isinstance(zh_li_term.u, numbers.Real):
                 u = df.Field(
                     mesh=system.m.mesh,
                     nvdim=3,
-                    value=(1.0, 0.0, 0.0),
-                    norm=zh_li_term.u,
+                    value=(zh_li_term.u, 0.0, 0.0),
                 )
             elif isinstance(zh_li_term.u, dict):
-                raise NotImplementedError(
-                    "Setting Zhang Li u with a 'dict' is not yet supported."
+                if isinstance(list(zh_li_term.u.values())[0], numbers.Real):
+                    u_values = {key: (value, 0.0, 0.0) for key, value in zh_li_term.u.items()}
+                else:
+                    u_values = zh_li_term.u
+                u = df.Field(
+                    mesh=system.m.mesh,
+                    nvdim=3,
+                    value=u_values,
                 )
             else:  # array_like
                 u = df.Field(mesh=system.m.mesh, nvdim=3, value=zh_li_term.u)
diff --git a/mumax3c/tests/conftest.py b/mumax3c/tests/conftest.py
index 75a07e2..966224e 100644
--- a/mumax3c/tests/conftest.py
+++ b/mumax3c/tests/conftest.py
@@ -32,7 +32,7 @@
     "TestDamping.test_field",
     "TestDynamics.test_field_field",
     "TestSlonczewski.test_field_values",
-    "TestZhangLi.test_field_scalar",
+    "TestZhangLi.test_time_tcl_scalar_u",
     "TestUniaxialAnisotropy.test_field_vector",
     "TestUniaxialAnisotropy.test_scalar_field",
     "TestUniaxialAnisotropy.test_field_field",
@@ -54,8 +54,7 @@
     "TestSlonczewski.test_single_values",
     "TestSlonczewski.test_single_values_finite_temperature",
     "TestSlonczewski.test_dict_values",
-    "TestZhangLi.test_dict_scalar",
-    "TestZhangLi.test_time_scalar_scalar",
+    "TestZhangLi.test_time_func_scalar_u",
     "test_multiple_drives_compute",
     "TestPrecession.test_scalar",
     "TestDMI.test_scalar",