Apply various easy fixes

Author: LourensVeen

src/amuse/examples/distinct_colours.py

@@ -0,0 +1,92 @@
+# -*- coding: iso-8859-1 -*-
+
+"""
+Colour-blind proof distinct colours module, based on work by Paul Tol
+Pieter van der Meer, 2011
+SRON - Netherlands Institute for Space Research
+"""
+
+# colour table in HTML hex format
+hexcols = ['#332288', '#88CCEE', '#44AA99', '#117733', '#999933', '#DDCC77', 
+           '#CC6677', '#882255', '#AA4499', '#661100', '#6699CC', '#AA4466',
+           '#4477AA']
+
+greysafecols = ['#809BC8', '#FF6666', '#FFCC66', '#64C204']
+
+xarr = [[12], 
+        [12, 6], 
+        [12, 6, 5], 
+        [12, 6, 5, 3], 
+        [0, 1, 3, 5, 6], 
+        [0, 1, 3, 5, 6, 8], 
+        [0, 1, 2, 3, 5, 6, 8], 
+        [0, 1, 2, 3, 4, 5, 6, 8], 
+        [0, 1, 2, 3, 4, 5, 6, 7, 8], 
+        [0, 1, 2, 3, 4, 5, 9, 6, 7, 8], 
+        [0, 10, 1, 2, 3, 4, 5, 9, 6, 7, 8], 
+        [0, 10, 1, 2, 3, 4, 5, 9, 6, 11, 7, 8]]
+
+# get specified nr of distinct colours in HTML hex format.
+# in: nr - number of colours [1..12]
+# returns: list of distinct colours in HTML hex
+
+def get_distinct(nr):
+
+    # check if nr is in correct range
+    
+    if nr < 1 or nr > 12:
+        print("wrong nr of distinct colours!")
+        return
+
+    # get list of indices
+    
+    lst = xarr[nr-1]
+    
+    # generate colour list by stepping through indices and looking them up
+    # in the colour table
+
+    i_col = 0
+    col = [0] * nr
+    for idx in lst:
+        col[i_col] = hexcols[idx]
+        i_col+=1
+    return col
+
+# gets 4 colours, which also look distinct in black&white
+# returns: list of 4 colours in 
+#def get_distinct_grey():
+
+import sys
+
+# display usage information and produce example plot.
+if __name__ == '__main__':
+    import numpy as np
+    import matplotlib.mlab as mlab
+    import matplotlib.pyplot as plt
+
+    print(__doc__)
+    print("usage examples: ")
+    print("print distinct_colours.get_distinct(2)")
+    print(get_distinct(2))
+    print("print distinct_colours.greysafecols")
+    print(greysafecols)
+
+    print("generating example plot: distinct_colours_example.png")
+    plt.close()
+    t = np.arange(0.0, 2.0, 0.01)
+    n = 12
+    if len(sys.argv) > 1: n = int(sys.argv[1])
+    cols = get_distinct(n)
+    d = 2./n
+    for i in range(n):
+        s = np.sin(i*d*np.pi*t)
+        plt.plot(t, s, linewidth=2.0, c=cols[i], label=str(i))
+
+    plt.xlabel('time (s)')
+    plt.ylabel('voltage (mV)')
+    plt.title('Distinct colours example (n = '+str(n)+')')
+    plt.grid(True)
+    plt.legend(loc=3, ncol=2)
+    
+    plt.savefig("distinct_colours_example.png")
+    plt.show()

src/amuse/examples/gravity_hydro.py

@@ -152,7 +152,12 @@ def new_option_parser():
                       help="end time of the simulation [%default]")
     return result
 
-if __name__ in ('__main__', '__plot__'):
+
+def  main():
     o, arguments  = new_option_parser().parse_args()
     numpy.random.seed(123)
     gravity_hydro_bridge(**o.__dict__)
+
+
+if __name__ in ('__main__', '__plot__'):
+    main()

src/amuse/examples/gravity_stellar_minimal.py

@@ -166,13 +166,13 @@ def gravity_stellar_minimal_argument_parser():
     return result
 
 
-def main(args):
-    arguments = gravity_stellar_minimal_argument_parser().parse_args(args)
-    gravity_stellar_minimal(**arguments.__dict__)
+def main():
+    if shell_is_interactive():
+        gravity_stellar_minimal({})
+    else:
+        arguments = gravity_stellar_minimal_argument_parser().parse_args(sys.argv[1:])
+        gravity_stellar_minimal(**arguments.__dict__)
 
 
 if __name__ == "__main__":
-    if shell_is_interactive():
-        main([])
-    else:
-        main(sys.argv[1:])
+    main()

src/amuse/examples/gravity_to_virial.py

@@ -3,8 +3,8 @@
 """
 from amuse.lab import *
 from matplotlib import pyplot
-from prepare_figure import single_frame, figure_frame, set_tickmarks
-from distinct_colours import get_distinct
+from amuse.examples.prepare_figure import single_frame, figure_frame, set_tickmarks
+from amuse.examples.distinct_colours import get_distinct
 
 def virial_ratio_evolution(code, bodies, Q_init, t_end):
     dt = 0.06125 | t_end.unit

src/amuse/examples/hydrodynamics_class.py

@@ -8,7 +8,7 @@
 from amuse.community.fi.interface import Fi
 from amuse.community.gadget2.interface import Gadget2
 
-from cooling_class import Cooling, SimplifiedThermalModelEvolver
+from amuse.examples.cooling_class import Cooling, SimplifiedThermalModelEvolver
 # from amuse.ext.sink import SinkParticles
 from amuse.ext.sink import new_sink_particles
 

src/amuse/examples/iliev_test5.py

@@ -344,11 +344,9 @@ def plots(i):
     aplot(i, 'mach', ((r/15, mach, 'r'),),
           xlim=(0., 1.), ylim=(1.e-5, 10.))
 
-# main example
-
-
-if __name__ == "__main__":
 
+# main example
+def main():
     N = 10000
     Ns = 1
     L = 15. | units.kpc
@@ -361,3 +359,7 @@ def plots(i):
     if IS_PLOT_AVAILABLE:
         plots(50)
         plots(100)
+
+
+if __name__ == "__main__":
+    main()

src/amuse/examples/kelvin_helmholtz.py

@@ -152,7 +152,12 @@ def plot_grid(grid):
     plot.imshow(rho, origin = 'lower', cmap='jet')
     figure.savefig('kelvin_helmholtz.png')
     pyplot.show()
-    
-if __name__ == "__main__":
+
+
+def main():
     grids = simulate_kelvin_helmholtz_instability(1.0 | time)
     plot_grid(grids[0])
+
+
+if __name__ == "__main__":
+    main()

src/amuse/examples/list_of_examples.py

@@ -5,7 +5,7 @@
 chapter_6 = ["hydro_minimal", "hydro_outflow_particles", "hydro_sink_particles", "hydro_disk_with_bump", "explode_evolved_star", "helium_star", ]
 chapter_7 = ["rad_minimal", "radiative_h2cloud", ]
 chapter_8 = ["gravity_stellar_minimal", "gravity_stellar_collision", "plot_fractal_clumpy_cluster", "gravity_stellar_eventdriven", "merge_two_stars_and_evolve", "merge_two_stars_sph", "gravity_kepler_disks", "basic_multiples", "kira", "lonely_planets", ]
-chapter_9 = ["gravity_potential", "gravity_drifter", "two_body_bridge", "three_body_bridge", "three_body_bridge_hierarchical", "gravity_hydro", "relas_gas_and_stars", "evolve_xi_tau_primary", ]
+chapter_9 = ["gravity_potential", "gravity_drifter", "two_body_bridge", "three_body_bridge", "three_body_bridge_hierarchical", "gravity_hydro", "relax_gas_and_stars", "evolve_xi_tau_primary", ]
 chapter_10 = ["multiples_minimal", "kira", ]  
 appendix_a = ["multiple_stellar_threaded", ]
 appendix_b = []

src/amuse/examples/make_oort_cloud.py

@@ -393,7 +393,7 @@ def main():
     outfile = args.outfile
 
     if args.outfile is None:
-        outfile = f"sun_and_comets_N{args:01}.amuse"
+        outfile = f"sun_and_comets_N{args.n_comets:01}.amuse"
 
     bodies = read_set_from_file(args.filename, close_file=True)
     stars = bodies[bodies.type == "star"]

src/amuse/examples/merge_two_stars_sph_convergence.py

@@ -4,8 +4,8 @@
 from amuse.lab import *
 from amuse.plot import plot
 from amuse.ext.sph_to_star import convert_SPH_to_stellar_model
-from prepare_figure import *
-from distinct_colours import get_distinct
+from amuse.examples.prepare_figure import *
+from amuse.examples.distinct_colours import get_distinct
 
 def return_evolved_star_hydro(mass, time, Nsph):
     star =  Particle(mass=mass)