Skip to content

Commit

Permalink
RFC: refactor benchmarks setups and customize pytest's discovery rule…
Browse files Browse the repository at this point in the history
… to enable running benchmarks through pytest
  • Loading branch information
neutrinoceros committed May 15, 2024
1 parent 81482cf commit 3bde28c
Show file tree
Hide file tree
Showing 14 changed files with 88 additions and 5 deletions.
9 changes: 9 additions & 0 deletions benchmarks/coordinates.py
Original file line number Diff line number Diff line change
Expand Up @@ -51,6 +51,9 @@ def setup(self):
self.array_rep = CartesianRepresentation(np.ones((3, 1000)) * u.kpc)
self.array_dif = CartesianDifferential(np.ones((3, 1000)) * u.km / u.s)

# pytest compat
setup_method = setup

def time_with_differentials_scalar(self):
self.scalar_rep.with_differentials(self.scalar_dif)

Expand Down Expand Up @@ -129,6 +132,9 @@ def setup(self):
xyz_clustered2, representation_type=CartesianRepresentation
)

# pytest compat
setup_method = setup

def time_init_nodata(self):
FK5()

Expand Down Expand Up @@ -188,6 +194,9 @@ def setup(self):
lat=self.array_q_dec, lon=self.array_q_ra
)

# pytest compat
setup_method = setup

def time_init_scalar(self):
SkyCoord(1, 2, unit="deg", frame="icrs")

Expand Down
3 changes: 3 additions & 0 deletions benchmarks/io_ascii/core.py
Original file line number Diff line number Diff line change
Expand Up @@ -41,6 +41,9 @@ def setup(self):
for col, x in izip(self.cols, lst):
col.str_vals = [str(s) for s in x]

# pytest compat
setup_method = setup

def time_continuation_inputter(self):
core.ContinuationLinesInputter().process_lines(self.lines)

Expand Down
3 changes: 3 additions & 0 deletions benchmarks/io_ascii/fixedwidth.py
Original file line number Diff line number Diff line change
Expand Up @@ -18,6 +18,9 @@ def setup(self):
self.splitter.cols = self.header.cols
self.data = ascii.FixedWidthData()

# pytest compat
setup_method = setup

def time_splitter(self):
self.splitter(self.lines[1:])

Expand Down
3 changes: 3 additions & 0 deletions benchmarks/io_ascii/ipac.py
Original file line number Diff line number Diff line change
Expand Up @@ -26,6 +26,9 @@ def setup(self):
self.data.cols = list(self.table.columns.values())
self.data._set_fill_values(self.data.cols)

# pytest compat
setup_method = setup

def time_splitter(self):
self.splitter.join(self.vals, self.widths)

Expand Down
3 changes: 3 additions & 0 deletions benchmarks/io_ascii/main.py
Original file line number Diff line number Diff line change
Expand Up @@ -43,6 +43,9 @@ def setup(self):
if self.file_format != "sextractor":
self.table = self.read()

# pytest compat
setup_method = setup

def read(self):
return ascii.read(BytesIO(self.data), format=self.file_format, guess=False)

Expand Down
3 changes: 3 additions & 0 deletions benchmarks/io_ascii/rdb.py
Original file line number Diff line number Diff line change
Expand Up @@ -13,5 +13,8 @@ def setup(self):
self.lines = f.read().split("\n")
f.close()

# pytest compat
setup_method = setup

def time_get_cols(self):
self.header.get_cols(self.lines)
3 changes: 3 additions & 0 deletions benchmarks/io_ascii/sextractor.py
Original file line number Diff line number Diff line change
Expand Up @@ -23,5 +23,8 @@ def setup(self):
self.lines.append("# {} {} Description [pixel**2]".format(i, randword()))
self.lines.append("Non-header line")

# pytest compat
setup_method = setup

def time_header(self):
self.header.get_cols(self.lines)
3 changes: 3 additions & 0 deletions benchmarks/io_ascii/table.py
Original file line number Diff line number Diff line change
Expand Up @@ -28,6 +28,9 @@ def setup(self):
self.outputter = core.TableOutputter()
self.table = table.Table()

# pytest compat
setup_method = setup

def time_table_outputter(self):
self.outputter(self.cols, {"table": {}})

Expand Down
18 changes: 13 additions & 5 deletions benchmarks/io_fits.py
Original file line number Diff line number Diff line change
Expand Up @@ -25,6 +25,9 @@ def setup(self):
t["booleans"] = t["floats"] > 0.5
t.write(temp, format="fits")

# pytest compat
setup_method = setup

def time_read_nommap(self):
try:
Table.read(self.table_file, format="fits", memmap=False)
Expand All @@ -42,7 +45,7 @@ def time_write(self):
t.write(table_bytes, format="fits")


class FITSBinTableHDU:
class FITSBinTableHDUBenchmarks:
def time_from_columns_bytes(self):
x = np.repeat(b"a", 2_000_000)
array = np.array(x, dtype=[("col", "S1")])
Expand All @@ -57,7 +60,7 @@ def make_header(ncards=1000):
return fits.Header(cards)


class FITSHeader:
class FITSHeaderBenchmarks:
"""
Tests of the Header interface
"""
Expand All @@ -66,6 +69,9 @@ def setup(self):
self.hdr = make_header()
self.hdr_string = self.hdr.tostring()

# pytest compat
setup_method = setup

def time_get_int(self):
self.hdr.get("INT999")

Expand All @@ -85,21 +91,23 @@ def time_fromstring(self):
fits.Header.fromstring(self.hdr_string)


class FITSHDUList:
class FITSHDUListBenchmarks:
"""
Tests of the HDUList interface
"""

filename = "many_hdu.fits"

def setup_cache(self):
self.filename = NamedTemporaryFile(delete=False).name
hdr = make_header()
hdul = fits.HDUList(
[fits.PrimaryHDU(header=hdr)]
+ [fits.ImageHDU(header=hdr) for _ in range(30)]
)
hdul.writeto(self.filename)

# pytest compat
setup_method = setup_cache

def time_getheader(self):
fits.getheader(self.filename)

Expand Down
6 changes: 6 additions & 0 deletions benchmarks/modeling/compound.py
Original file line number Diff line number Diff line change
Expand Up @@ -47,6 +47,9 @@ def setup(self):
| models.RotateNative2Celestial(5.6, -72.05, 180)
)

# pytest compat
setup_method = setup

def time_scalar(self):
r, d = self.model(x_no_units_scalar, x_no_units_scalar)

Expand Down Expand Up @@ -77,6 +80,9 @@ def setup(self):
| models.RotateNative2Celestial(5.6 * u.deg, -72.05 * u.deg, 180 * u.deg)
)

# pytest compat
setup_method = setup

def time_scalar(self):
r, d = self.model(x_no_units_scalar * u.pix, x_no_units_scalar * u.pix)

Expand Down
3 changes: 3 additions & 0 deletions benchmarks/stats/sigma_clipping.py
Original file line number Diff line number Diff line change
Expand Up @@ -26,6 +26,9 @@ def setup(self):
# deviation as the stdfunc. The default iters is 5.
self.sigclip = SigmaClip(sigma=3)

# pytest compat
setup_method = setup

def time_3d_array(self):
self.sigclip(self.data[:, :1024, :1024])

Expand Down
12 changes: 12 additions & 0 deletions benchmarks/table.py
Original file line number Diff line number Diff line change
Expand Up @@ -42,6 +42,9 @@ def setup(self):

self.bool_mask = self.table["a"] > 0.6

# pytest compat
setup_method = setup

def time_table_slice_bool(self):
self.table[self.bool_mask]

Expand Down Expand Up @@ -148,6 +151,9 @@ class TimeMaskedColumn:
def setup(self):
self.dat = np.arange(1e7)

# pytest compat
setup_method = setup

def time_masked_column_init(self):
MaskedColumn(self.dat)

Expand All @@ -156,6 +162,9 @@ class TimeTableInitWithLists:
def setup(self):
self.dat = list(range(100_000))

# pytest compat
setup_method = setup

def time_init_lists(self):
Table([self.dat, self.dat, self.dat], names=["time", "rate", "error"])

Expand Down Expand Up @@ -183,6 +192,9 @@ def setup(self):
self.data_str_masked_1d = self.data_str_1d.copy()
self.data_str_masked_1d[-1] = np.ma.masked

# pytest compat
setup_method = setup

def time_init_int_1d(self):
Table([self.data_int_1d])

Expand Down
15 changes: 15 additions & 0 deletions benchmarks/units.py
Original file line number Diff line number Diff line change
Expand Up @@ -119,6 +119,9 @@ def setup(self):
self.out_sq = data * u.g**2
self.out_sqrt = data * u.g**0.5

# pytest compat
setup_method = setup

def time_quantity_square(self):
self.data**2

Expand Down Expand Up @@ -149,6 +152,8 @@ def setup(self):
self.out_sq = data * u.g**2
self.out_sqrt = data * u.g**0.5

# pytest compat
setup_method = setup

class TimeQuantityOpSmallArrayDiffUnit:
"""
Expand All @@ -163,6 +168,9 @@ def setup(self):
# A different but dimensionally compatible unit
self.data2 = 0.001 * data * u.kg

# pytest compat
setup_method = setup

def time_quantity_equal(self):
# Same as operator.eq
self.data == self.data2
Expand Down Expand Up @@ -211,6 +219,8 @@ def setup(self):
self.data = data * u.g
self.data2 = self.data.copy()

# pytest compat
setup_method = setup

class TimeQuantityOpLargeArrayDiffUnit(TimeQuantityOpSmallArrayDiffUnit):
"""
Expand All @@ -224,6 +234,8 @@ def setup(self):
# A different but dimensionally compatible unit
self.data2 = 0.001 * data * u.kg

# pytest compat
setup_method = setup

class TimeQuantityOpLargeArraySameUnit(TimeQuantityOpSmallArrayDiffUnit):
"""
Expand All @@ -234,3 +246,6 @@ def setup(self):
data = np.arange(1e6) + 1
self.data = data * u.g
self.data2 = self.data.copy()

# pytest compat
setup_method = setup
9 changes: 9 additions & 0 deletions pytest.ini
Original file line number Diff line number Diff line change
@@ -0,0 +1,9 @@
[pytest]
# customize test discovery to treat benchmarks as tests
python_files = *.py
python_classes = Time *Benchmarks
python_functions = time_*, timeraw_*, mem_*, peakmem_*, and track_*
addopts =
# ignored because pytest interprets some method arguments as missing fixtures
--ignore=benchmarks/cosmology.py
--ignore=benchmarks/convolve.py

0 comments on commit 3bde28c

Please sign in to comment.