diff --git a/pyproject.toml b/pyproject.toml index df1e4247..fb4a7eb6 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -23,6 +23,7 @@ dependencies = [ "numpy>=2.0,<2.5", "click (>=8.2.1,<9.0.0)", "rustworkx>=0.17.1,<1.0", + "matplotlib>=3.10.8", ] [dependency-groups] diff --git a/python/opencosmo/analysis/plots.py b/python/opencosmo/analysis/plots.py new file mode 100644 index 00000000..29d80cac --- /dev/null +++ b/python/opencosmo/analysis/plots.py @@ -0,0 +1,82 @@ +import opencosmo as oc +from opencosmo.analysis import reduce +import matplotlib.pyplot as plt +import numpy as np + +def hist1d(data, column_name, bins = 20): + + quantity = data#ds.select(column_name).get_data("numpy") + bin_edges = np.logspace(np.log10(quantity.min()), np.log10(quantity.max()), bins) + + plt.hist(quantity, bins = bin_edges, edgecolor='black') + plt.xlabel(column_name) + plt.ylabel('Counts') + plt.yscale('log') + plt.xscale('log') + plt.savefig(f'{column_name}_hist.png') + plt.show() + plt.close() + + return + +def hist2d(ds, x_column, y_column, gridsize = 50, cmap = 'Blues'): + + x = ds.select(x_column).get_data("numpy") + y = ds.select(y_column).get_data("numpy") + + hb = plt.hexbin(x, y, gridsize=gridsize, bins='log', xscale = 'log', yscale = 'log', cmap=cmap) + cb = plt.colorbar(hb, label='Counts') + + plt.xlabel(x_column) + plt.ylabel(y_column) + plt.show() + + return + +def scatter(ds, x_column, y_column, xscale = 'log', yscale = 'log', labels = None, **kwargs): + + x = ds.select(x_column).get_data("numpy") + y = ds.select(y_column).get_data("numpy") + + for k, val in kwargs.items(): + plt.scatter(x, y, **kwargs) + + if labels == None: + plt.xlabel(x_column) + plt.ylabel(y_column) + + else: + plt.xlabel(labels[0]) + plt.ylabel(labels[1]) + + plt.xscale(xscale) + plt.yscale(yscale) + plt.show() + + return + +def reduce_hmf(sod_halo_mass, bins, boxsize): + b = np.logspace(np.log10(sod_halo_mass).min(), np.log10(sod_halo_mass).max(), bins) + h, bin_edge = np.histogram(sod_halo_mass, bins = b) + + mbins = 0.5 * (bin_edge[:-1] + bin_edge[1:]) + hmf = np.abs(h/np.diff(bin_edge)/boxsize**3) + err = np.abs(np.sqrt(h)/np.diff(bin_edge)/boxsize**3) + # log_err = err / (np.log(10)*hmf) + + plt.plot(mbins, hmf, linestyle = '--', color = 'k') + plt.errorbar(mbins, hmf, yerr = err , color = 'k', fmt = 'o') + plt.xscale('log') + plt.yscale('log') + plt.xlabel(r'Halo Mass [$M_\odot$]') + plt.ylabel(r'Number Density $\frac{dn}{dM}$') + plt.show() + plt.close() + +def HMF(ds, bins, boxsize): + evaluate_kwargs = {"format": "numpy", "vectorize" : True, "insert" : False, "boxsize": boxsize, "bins": bins} + reduce(ds, reduce_hmf, evaluate_kwargs = evaluate_kwargs) + + + + \ No newline at end of file diff --git a/python/opencosmo/analysis/test.py b/python/opencosmo/analysis/test.py new file mode 100644 index 00000000..632da737 --- /dev/null +++ b/python/opencosmo/analysis/test.py @@ -0,0 +1,17 @@ +import opencosmo as oc +from opencosmo.analysis import reduce +from opencosmo.analysis.plots import hist1d, hist2d, scatter, HMF + +ds = oc.open('python/opencosmo/analysis/haloproperties.hdf5').filter(oc.col('sod_halo_mass') > 0).take(200) + +# hist1d(ds,"sod_halo_mass") + +# hist2d(ds, "sod_halo_mass", "fof_halo_mass") + +# #set s=0.1 to test kwargs +# scatter(ds, "sod_halo_mass", "fof_halo_mass", s = 0.1) + +# evaluate_kwargs = {"format": "numpy", "vectorize" : True, "insert" : False, "boxsize": 10, "bins": 15} +# reduce(ds, HMF, evaluate_kwargs = evaluate_kwargs) + +HMF(ds, bins = 15, boxsize = 10) \ No newline at end of file diff --git a/uv.lock b/uv.lock index 169948f5..987203f0 100644 --- a/uv.lock +++ b/uv.lock @@ -1035,6 +1035,7 @@ dependencies = [ { name = "hdf5plugin" }, { name = "healpy" }, { name = "healsparse" }, + { name = "matplotlib" }, { name = "numpy" }, { name = "pydantic" }, { name = "rustworkx" }, @@ -1085,6 +1086,7 @@ requires-dist = [ { name = "hdf5plugin", specifier = ">=5.0.0,!=5.1,<7.0" }, { name = "healpy", specifier = ">=1.19.0,<2.0.0" }, { name = "healsparse", specifier = ">=1.11,<2.0" }, + { name = "matplotlib", specifier = ">=3.10.8" }, { name = "numpy", specifier = ">=2.0,<2.5" }, { name = "pyarrow", marker = "extra == 'io'", specifier = ">=21.0.0" }, { name = "pydantic", specifier = ">=2.10.6,<3.0.0" },