From 6e4d52f4ed3dc20d2787bb40f2a4ef5b04adf584 Mon Sep 17 00:00:00 2001 From: Michael Coughlin Date: Tue, 7 May 2024 14:35:40 -0500 Subject: [PATCH] start demoing MILP --- gwemopt/args.py | 4 ++ gwemopt/params.py | 11 +++ gwemopt/scheduler.py | 141 ++++++++++++++++++++++++++++++++++++-- gwemopt/utils/__init__.py | 1 + gwemopt/utils/milp.py | 125 +++++++++++++++++++++++++++++++++ pyproject.toml | 1 + 6 files changed, 276 insertions(+), 7 deletions(-) create mode 100644 gwemopt/utils/milp.py diff --git a/gwemopt/args.py b/gwemopt/args.py index 6aceb9b..e4f1d3c 100644 --- a/gwemopt/args.py +++ b/gwemopt/args.py @@ -160,4 +160,8 @@ def parse_args(args): parser.add_argument("--inclination", action="store_true", default=False) + parser.add_argument("--solverType", default="heuristic") + parser.add_argument("--milpSolver", default="PULP_CBC_CMD") + parser.add_argument("--milpOptions") + return parser.parse_args(args=args) diff --git a/gwemopt/params.py b/gwemopt/params.py index 07211a2..b0c9f9b 100644 --- a/gwemopt/params.py +++ b/gwemopt/params.py @@ -1,3 +1,4 @@ +import json import os from pathlib import Path @@ -146,4 +147,14 @@ def params_struct(opts): opts.inclination if hasattr(opts, "true_location") else False ) + params["solverType"] = ( + opts.solverType if hasattr(opts, "solverType") else "heuristic" + ) + params["milpSolver"] = ( + opts.milpSolver if hasattr(opts, "milpSolver") else "PULP_CBC_CMD" + ) + params["milpOptions"] = ( + json.loads(opts.milpOptions) if hasattr(opts, "milpOptions") else {} + ) + return params diff --git a/gwemopt/scheduler.py b/gwemopt/scheduler.py index a5dc684..3ec9efa 100644 --- a/gwemopt/scheduler.py +++ b/gwemopt/scheduler.py @@ -9,9 +9,7 @@ from ortools.linear_solver import pywraplp from gwemopt.tiles import balance_tiles, optimize_max_tiles, schedule_alternating -from gwemopt.utils import angular_distance - -# from munkres import Munkres, make_cost_matrix +from gwemopt.utils import angular_distance, solve_milp def get_altaz_tile(ra, dec, observer, obstime): @@ -90,7 +88,7 @@ def find_tile( return idx2, exposureids, probs -def get_order( +def get_order_heuristic( params, tile_struct, tilesegmentlists, exposurelist, observer, config_struct ): """ @@ -473,6 +471,127 @@ def get_order( return idxs, filts +def get_order_milp(params, tile_struct, exposurelist, observer, config_struct): + """ + tile_struct: dictionary. key -> struct info. + exposurelist: list of segments that the telescope is supposed to be working. + consecutive segments from the start to the end, with each segment size + being the exposure time. + Returns a list of tile indices in the order of observation. + """ + + if "dec_constraint" in config_struct: + dec_constraint = config_struct["dec_constraint"].split(",") + dec_min = float(dec_constraint[0]) + dec_max = float(dec_constraint[1]) + + exposureids = [] + probs = [] + ras, decs, filts, keys = [], [], [], [] + for ii, key in enumerate(list(tile_struct.keys())): + if tile_struct[key]["prob"] == 0: + continue + if "dec_constraint" in config_struct: + if (tile_struct[key]["dec"] < dec_min) or ( + tile_struct[key]["dec"] > dec_max + ): + continue + + exposureids.append(key) + probs.append(tile_struct[key]["prob"]) + ras.append(tile_struct[key]["ra"]) + decs.append(tile_struct[key]["dec"]) + filts.append(tile_struct[key]["filt"]) + keys.append(key) + + fields = -1 * np.ones( + (len(exposurelist)), + ) + filters = ["n"] * len(exposurelist) + + if len(probs) == 0: + return fields, filters + + probs = np.array(probs) + ras = np.array(ras) + decs = np.array(decs) + exposureids = np.array(exposureids) + keys = np.array(keys) + tilematrix = np.zeros((len(exposurelist), len(ras))) + for ii in np.arange(len(exposurelist)): + # first, create an array of airmass-weighted probabilities + t = Time(exposurelist[ii][0], format="mjd") + altazs = [get_altaz_tile(ra, dec, observer, t) for ra, dec in zip(ras, decs)] + alts = np.array([altaz[0] for altaz in altazs]) + horizon = config_struct["horizon"] + horizon_mask = alts <= horizon + airmass = 1 / np.cos((90.0 - alts) * np.pi / 180.0) + airmass_mask = airmass > params["airmass"] + + airmass_weight = 10 ** (0.4 * 0.1 * (airmass - 1)) + + if params["scheduleType"] in ["greedy", "greedy_slew"]: + tilematrix[ii, :] = np.array(probs) + elif params["scheduleType"] == ["airmass_weighted", "airmass_weighted_slew"]: + tilematrix[ii, :] = np.array(probs / airmass_weight) + tilematrix[ii, horizon_mask] = np.nan + tilematrix[ii, airmass_mask] = np.nan + + for jj, key in enumerate(keys): + tilesegmentlist = tile_struct[key]["segmentlist"] + if not tilesegmentlist.intersects_segment(exposurelist[ii]): + tilematrix[ii, jj] = np.nan + + # which fields are never observable + ind = np.where(np.nansum(tilematrix, axis=0) > 0)[0] + + probs = np.array(probs)[ind] + ras = np.array(ras)[ind] + decs = np.array(decs)[ind] + exposureids = np.array(exposureids)[ind] + filts = [filts[i] for i in ind] + tilematrix = tilematrix[:, ind] + + # which times do not have any observability + ind = np.where(np.nansum(tilematrix, axis=1) > 0)[0] + tilematrix = tilematrix[ind, :] + + cost_matrix = tilematrix + cost_matrix[np.isnan(cost_matrix)] = -np.inf + + distmatrix = np.zeros((len(ras), len(ras))) + for ii, (r, d) in enumerate(zip(ras, decs)): + dist = angular_distance(r, d, ras, decs) + if "slew" in params["scheduleType"]: + dist = dist / config_struct["slew_rate"] + dist = dist - config_struct["readout"] + dist[dist < 0] = 0 + else: + distmatrix[ii, :] = dist + distmatrix = distmatrix / np.max(distmatrix) + + dt = int(np.ceil((exposurelist[1][0] - exposurelist[0][0]) * 86400)) + optimal_points = solve_milp( + cost_matrix, + dist_matrix=distmatrix, + useDistance=False, + max_tasks_per_worker=len(params["filters"]), + useTaskSepration=False, + min_task_separation=int(np.ceil(dt / params["mindiff"])), + ) + + for optimal_point in optimal_points: + idx = ind[optimal_point[0]] + idy = optimal_point[1] + if len(filts[idy]) > 0: + fields[idx] = exposureids[idy] + filters[idx] = filts[idy][0] + filt = filts[idy][1:] + filts[idy] = filt + + return fields, filters + + def scheduler(params, config_struct, tile_struct): """ config_struct: the telescope configurations @@ -506,9 +625,17 @@ def scheduler(params, config_struct, tile_struct): for key in keys: # segments.py: tile_struct[key]["segmentlist"] is a list of segments when the tile is available for observation tilesegmentlists.append(tile_struct[key]["segmentlist"]) - keys, filts = get_order( - params, tile_struct, tilesegmentlists, exposurelist, observer, config_struct - ) + + if params["solverType"] == "heuristic": + keys, filts = get_order_heuristic( + params, tile_struct, tilesegmentlists, exposurelist, observer, config_struct + ) + elif params["solverType"] == "milp": + keys, filts = get_order_milp( + params, tile_struct, exposurelist, observer, config_struct + ) + else: + raise ValueError(f'Unknown solverType {params["solverType"]}') if params["doPlots"]: from gwemopt.plotting import make_schedule_plots diff --git a/gwemopt/utils/__init__.py b/gwemopt/utils/__init__.py index 9fbf2df..033e505 100644 --- a/gwemopt/utils/__init__.py +++ b/gwemopt/utils/__init__.py @@ -1,4 +1,5 @@ from gwemopt.utils.geometry import angular_distance +from gwemopt.utils.milp import solve_milp from gwemopt.utils.misc import auto_rasplit, get_exposures, integrationTime from gwemopt.utils.observability import calculate_observability from gwemopt.utils.param_utils import readParamsFromFile diff --git a/gwemopt/utils/milp.py b/gwemopt/utils/milp.py new file mode 100644 index 0000000..7ee0022 --- /dev/null +++ b/gwemopt/utils/milp.py @@ -0,0 +1,125 @@ +import time + +import numpy as np +import pulp +from tqdm import tqdm + + +def solve_milp( + cost_matrix, + max_tasks_per_worker=1, + useTaskSepration=False, + min_task_separation=1, + useDistance=False, + dist_matrix=None, + timeLimit=300, + milpSolver="PULP_CBC_CMD", + milpOptions={}, +): + + cost_matrix_mask = cost_matrix > 10 ** (-10) + optimal_points = [] + + if cost_matrix_mask.any(): + print("Calculating MILP solution...") + + # Create a CP-SAT model + problem = pulp.LpProblem("problem", pulp.LpMaximize) + + # Define variables + num_exposures, num_fields = cost_matrix.shape + + print("Define decision variables...") + # Define decision variables + x = { + (i, j): pulp.LpVariable(f"x_{i}_{j}", cat=pulp.LpBinary) + for i in tqdm(range(num_exposures)) + for j in range(num_fields) + } + + print("Define binary variables for task separation violation...") + # Define binary variables for task separation violation + s = { + (i, j): pulp.LpVariable(f"s_{i}_{j}", cat=pulp.LpBinary) + for i in tqdm(range(num_exposures)) + for j in range(num_fields) + } + + obj = pulp.lpSum( + x[i, j] * cost_matrix[i][j] + for i in range(num_exposures) + for j in range(num_fields) + ) + + if useDistance: + products = [ + (x[i, j], dist_matrix[j, k]) + for i in range(num_exposures) + for j in range(num_fields) + for k in range(dist_matrix.shape[1]) + ] + total_distance = pulp.LpAffineExpression(products) + obj -= total_distance + + # One field per exposure + for i in range(num_exposures): + problem += pulp.lpSum(x[i, j] for j in range(num_fields)) == 1 + + # Limit the number of tasks each worker can handle (if applicable) + for j in range(num_fields): + problem += ( + pulp.lpSum(x[i, j] for i in range(num_exposures)) + <= max_tasks_per_worker + ) + + print("Add constraints to exclude impossible assignments...") + # Add constraints to exclude impossible assignments + for i in tqdm(range(num_exposures)): + for j in range(num_fields): + if not np.isfinite(cost_matrix[i][j]): + problem += x[i, j] == 0 + + if useTaskSepration: + print("Define constraints: enforce minimum task separation...") + ## Define constraints: enforce minimum task separation + for i in tqdm(range(num_exposures)): + for j in range(num_fields): + for k in range(j + 1, num_fields): + problem += s[i, j] >= x[i, k] - x[i, j] - min_task_separation + problem += s[i, j] >= x[i, j] - x[i, k] - min_task_separation + + print("Number of variables:", len(problem.variables())) + print("Number of constraints:", len(problem.constraints)) + + time_limit = 60 # Stop the solver after 60 seconds + + if milpSolver == "PULP_CBC_CMD": + solver = pulp.getSolver(milpSolver) + elif milpSolver == "GUROBI": + solver = pulp.getSolver( + "GUROBI", manageEnv=True, envOptions=milpOptions, mip_gap=1e-6 + ) + else: + raise ValueError("milpSolver must be either PULP_CBC_CMD or GUROBI") + + solver.timeLimit = timeLimit + # solver.msg = True + status = problem.solve(solver) + + optimal_points = [] + if status in [pulp.LpStatusOptimal]: + for i in range(num_exposures): + for j in range(num_fields): + if ( + pulp.value(x[i, j]) is not None + and pulp.value(x[i, j]) > 0.5 + and np.isfinite(cost_matrix[i][j]) + ): + optimal_points.append((i, j)) + else: + print("The problem does not have a solution.") + + else: + print("The localization is not visible from the site.") + + return optimal_points diff --git a/pyproject.toml b/pyproject.toml index ae261b4..2ded42f 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -49,6 +49,7 @@ dependencies = [ 'healpy', 'ortools', 'pandas', + 'pulp', 'shapely', 'tables', 'regions'