Add a script for locating A_max

2 files changed, 73 insertions, 0 deletions

diff --git a/README b/README
index e41d2d5..a9333d4 100644
--- a/README
+++ b/README
@@ -39,6 +39,10 @@ public class that takes the C context options as keyword arguments and solves
 the equation on initialization. The instance methods then allow to evaluate the
 solution where needed.
 
+The following python scripts are bundled:
+- find_bound.py
+  Locates A_max using bisection
+
 Licence
 =======
 
diff --git a/find_bound.py b/find_bound.py
new file mode 100755
index 0000000..40ec22d
--- /dev/null
+++ b/find_bound.py
@@ -0,0 +1,69 @@
+#!/usr/bin/python3
+
+import argparse
+import sys
+
+import numpy as np
+
+import teukolsky_data
+
+parser = argparse.ArgumentParser('Find A_max for Teukolsky-type data')
+
+parser.add_argument('--coeffs-radial',  default = 64)
+parser.add_argument('--coeffs-angular', default = 16)
+parser.add_argument('--bounds-max',       type = float, default = np.nan)
+parser.add_argument('--bounds-min',       type = float, default = np.nan)
+parser.add_argument('--bounds-tolerance', type = float, default = 1e-8)
+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', default = 32)
+
+args = parser.parse_args()
+
+bounds    = [args.bounds_min, args.bounds_max]
+nb_coeffs = [args.coeffs_radial, args.coeffs_angular]
+
+if not (np.isnan(bounds[0]) or np.isnan(bounds[1])):
+    A = 0.5 * (bounds[0] + bounds[1])
+else:
+    A = 1.0
+
+def converged(bounds, tol):
+    return (np.isfinite(bounds[0]) and np.isfinite(bounds[1]) and
+            np.abs(bounds[0] - bounds[1]) < tol)
+
+sys.stderr.write('Solving with %dx%d coeffs\n' % (nb_coeffs[0], nb_coeffs[1]))
+
+prev_coeffs = None
+while not converged(bounds, args.bounds_tolerance):
+    sys.stderr.write('Trying A=%g (%s)\n' % (A, bounds))
+    try:
+        td_args = {
+            'amplitude'             : A,
+            'nb_coeffs'             : nb_coeffs,
+            'basis_scale_factor'    : [args.scale_factor, args.scale_factor],
+            'atol'                  : args.solve_tolerance,
+            'max_iter'              : args.max_iter,
+            'family'                : 0,
+        }
+
+        #if prev_coeffs:
+        #    args['coeffs_init'] = prev_coeffs
+
+        td = teukolsky_data.TeukolskyData(**td_args)
+
+        bounds[0] = A
+        sys.stderr.write('A=%16.16g converge\n' % A)
+        prev_coeffs = [c[:] for c in td.coeffs]
+    except Exception as e:
+        bounds[1] = A
+        sys.stderr.write('A=%16.16g diverge\n' % A)
+
+    if np.isnan(bounds[0]):
+        A *= 0.5
+    elif np.isnan(bounds[1]):
+        A *= 2
+    else:
+        A = 0.5 * (bounds[1] + bounds[0])
+
+sys.stdout.write('%16.16g %16.16g\n' % (bounds[0], bounds[1]))