From f9746f8e22ee33525abde628dfb37ed94f536b9e Mon Sep 17 00:00:00 2001
From: Anton Khirnov <anton@khirnov.net>
Date: Thu, 19 Jan 2023 15:48:30 +0100
Subject: Add a sample script for using the solver from Python

---
 README       |  2 ++
 teukolsky.py | 98 ++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++
 2 files changed, 100 insertions(+)
 create mode 100755 teukolsky.py

diff --git a/README b/README
index a9333d4..6e87fcb 100644
--- a/README
+++ b/README
@@ -40,6 +40,8 @@ the equation on initialization. The instance methods then allow to evaluate the
 solution where needed.
 
 The following python scripts are bundled:
+- teukolsky.py
+  Solves the constraints and prints the residuald and spectral coefficients
 - find_bound.py
   Locates A_max using bisection
 
diff --git a/teukolsky.py b/teukolsky.py
new file mode 100755
index 0000000..cde0c8f
--- /dev/null
+++ b/teukolsky.py
@@ -0,0 +1,98 @@
+#!/usr/bin/python3
+
+import argparse
+import numpy as np
+import sys
+
+import teukolsky_data
+
+parser = argparse.ArgumentParser('Solve the constraints for Teukolsky-type data')
+
+parser.add_argument('--coeffs-radial',    type = int,   default = 64)
+parser.add_argument('--coeffs-angular',   type = int,   default = 16)
+parser.add_argument('--solve-tolerance',  type = float, default = 1e-12)
+parser.add_argument('--scale-factor',     type = float, default = 3.0)
+parser.add_argument('--max-iter',         type = int,   default = 32)
+parser.add_argument('--family',           type = int,   default = 0)
+parser.add_argument('--upper-branch',     action = 'store_true')
+parser.add_argument('--print-coeffs',     action = 'store_true')
+
+parser.add_argument('amplitude', type = float)
+
+args = parser.parse_args()
+
+nb_coeffs       = [args.coeffs_radial, args.coeffs_angular]
+scale_factor    = [args.scale_factor, args.scale_factor]
+
+td = teukolsky_data.TeukolskyData(family = args.family, amplitude = args.amplitude,
+                                  nb_coeffs = nb_coeffs, basis_scale_factor = scale_factor,
+                                  atol = args.solve_tolerance, max_iter = args.max_iter,
+                                  solution_branch = int(args.upper_branch))
+
+r     = np.linspace(1e-3, 20, 256)
+theta = np.linspace(0, 2.0 * np.pi, 32)
+
+res = np.abs(td.eval_residual(r, theta))
+sys.stderr.write('residual: %g %g %g\n' % (np.max(res[0]), np.max(res[1]), np.max(res[2])))
+
+if args.print_coeffs:
+    for c in td.coeffs:
+        for j in range(nb_coeffs[1]):
+            for i in range(nb_coeffs[0]):
+                sys.stdout.write('%g ' % c[j, i])
+            sys.stdout.write('\n')
+
+        sys.stdout.write('\n')
+
+sys.exit(0)
+
+import matplotlib
+matplotlib.use('Agg')
+matplotlib.rc('text', usetex = True)
+matplotlib.rc('font', size = 8)
+matplotlib.rc('text.latex', preamble = r'\usepackage{times}\usepackage{mathrsfs}')
+matplotlib.rcParams['font.family'] = 'serif'
+
+from matplotlib import pyplot
+
+r     = np.linspace(0, 20, 1000)
+theta = np.linspace(0, 2.0 * np.pi, 200)
+R, Theta = np.meshgrid(r, theta)
+X = R * np.sin(Theta)
+Z = R * np.cos(Theta)
+
+varnames = ['$u$', r'$K_{r}^r$', r'$K_{\varphi}^\varphi$']
+vals = [td.eval_psi(r, theta) - 1.0, td.eval_krr(r, theta), td.eval_kpp(r, theta)]
+
+ax_lims = [[0, 1.0], [-1, 2], [-1, 1]]
+
+fig  = pyplot.figure(figsize = (6, 4))
+axes = [[None] * 2] * 3
+for i in range(3):
+    data = vals[i]
+    maxval = np.max(np.abs(data))
+
+    for j in range(2):
+        ax = fig.add_subplot(2, 3, j * 3 + i + 1)
+        axes[i][j] = ax
+
+        if j == 0:
+            ax.pcolormesh(X, Z, data, vmin = -maxval, vmax = maxval)
+            ax.set_xlim(-4, 4)
+            ax.set_ylim(-4, 4)
+            ax.set_aspect('equal')
+            ax.set_title(varnames[i])
+            ax.set_xlabel('$x$')
+            ax.set_ylabel('$z$')
+        else:
+            ax.plot(r,  data[0], 'b--')
+            ax.plot(-r, data[100], 'b--')
+
+            ax.plot(r,  data[50], 'g-')
+            ax.plot(-r, data[150], 'g-')
+
+            ax.set_xlim(-8, 8)
+            #ax.set_ylim(*ax_lims[i])
+
+fig.canvas.draw()
+fig.savefig(sys.argv[1], dpi = 600)
-- 
cgit v1.2.3