diff --git a/doc/energy-resource-build.nblink b/doc/energy-resource-build.nblink new file mode 100644 index 00000000..000de558 --- /dev/null +++ b/doc/energy-resource-build.nblink @@ -0,0 +1,3 @@ +{ + "path": "../examples/energy-resource-build.ipynb" +} diff --git a/doc/index.rst b/doc/index.rst index 39846607..41558585 100644 --- a/doc/index.rst +++ b/doc/index.rst @@ -120,6 +120,8 @@ This package is published under MIT license. :caption: Examples transport-tutorial + energy-resource-build + sudoku migrating-from-pyomo .. toctree:: diff --git a/doc/sudoku.nblink b/doc/sudoku.nblink new file mode 100644 index 00000000..5e6c9ac0 --- /dev/null +++ b/doc/sudoku.nblink @@ -0,0 +1,3 @@ +{ + "path": "../examples/sudoku.ipynb" +} diff --git a/doc/user-guide.rst b/doc/user-guide.rst index 8b7ee5bd..400545e1 100644 --- a/doc/user-guide.rst +++ b/doc/user-guide.rst @@ -48,7 +48,9 @@ Where to go next ---------------- - **Examples** — end-to-end problem walkthroughs: - :doc:`transport-tutorial`, :doc:`migrating-from-pyomo`. + :doc:`transport-tutorial`, :doc:`energy-resource-build` (power + portfolio build-out), :doc:`sudoku` (indexing with N-D coordinates), + :doc:`migrating-from-pyomo`. - **Advanced features** — :doc:`sos-constraints`, :doc:`piecewise-linear-constraints`, and the :doc:`testing-framework` for asserting structural properties of a diff --git a/examples/energy-resource-build.ipynb b/examples/energy-resource-build.ipynb new file mode 100644 index 00000000..acfc0362 --- /dev/null +++ b/examples/energy-resource-build.ipynb @@ -0,0 +1,445 @@ +{ + "cells": [ + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "# Power Portfolio Build-Out Example with Linopy\n", + "\n", + "This notebook provides a simplified example of building a power generation portfolio using `linopy`. The goal is to demonstrate the process of setting up and solving a linear optimization problem where different power generation resources are selected and built to meet future energy demands. " + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "import matplotlib.pyplot as plt\n", + "import numpy as np\n", + "import pandas as pd\n", + "import xarray as xr\n", + "\n", + "import linopy as lp" + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "## Problem Statement\n", + "This notebook presents an example portfolio expansion model used to meet an energy need. This was originally developed for a hydropower utility. Hydropower utilities can be energy limited (i.e., monthly energy balance) rather than capacity limited (i.e., hourly generating capacity), therefore, this algorithm addresses a monthly energy need rather than an hourly capacity need. A different algorithm would be needed to meet capacity needs and would include battery storage, etc. \n", + "\n", + "This notebook addresses the question: what is the least-cost way to meet an energy need in the future? " + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "## Generate Dummy Data\n", + "\n", + "A lot of this code is here just to generate dummy data. I would not spend too much time on this code. The `Linopy` model is down below. \n", + "### Energy Need\n", + "\n", + "This analysis assumes that an energy need has been identified through modeling and simulation. For this example, we will assume the energy need is as follows: " + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "def generate_dummy_energy_needs(start_year=2024, years=10):\n", + " # Generate a monthly date range\n", + " dates = pd.date_range(start=f\"{start_year}-01-01\", periods=years * 12, freq=\"M\")\n", + "\n", + " # Create a sinusoidal pattern for energy needs with a downward trend over time\n", + " months = np.arange(len(dates))\n", + "\n", + " # Base seasonal pattern using sine functions (with phase shift for winter, spring, etc.)\n", + " seasonal_pattern = np.sin(\n", + " 2 * np.pi * (months % 12) / 12\n", + " ) # Basic seasonal pattern (positive in spring/fall, negative in winter/summer)\n", + "\n", + " # Adjust the pattern so that winter months trend more negative over time\n", + " winter_weight = 0.5 * (\n", + " np.cos(2 * np.pi * (months % 12) / 12 - np.pi) + 1\n", + " ) # Emphasize winter more\n", + "\n", + " # Add a downward trend over time\n", + " trend = -0.01 * months # A small negative trend each month\n", + "\n", + " # Combine seasonal pattern, winter weight, and trend\n", + " energy_balance = seasonal_pattern * (1 - winter_weight) + trend\n", + " energy_balance *= 10\n", + " # Convert to pandas DataFrame\n", + " df = pd.DataFrame(data={\"monthly_energy_balance\": energy_balance}, index=dates)\n", + "\n", + " return df" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "# Generate the dummy dataset\n", + "dummy_energy_needs_df = generate_dummy_energy_needs()\n", + "dummy_energy_needs_df.plot()\n", + "plt.grid()" + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "The energy need is based on the negative monthly energy balance. This will be an input to the model, so it needs to be put in an xarray. " + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "dummy_energy_needs_df.index.name = \"datetime\"\n", + "dummy_energy_needs_df[\"energy_needs_aMW\"] = (\n", + " -dummy_energy_needs_df.monthly_energy_balance\n", + ").clip(lower=0)\n", + "energy_needs_xr = xr.DataArray(dummy_energy_needs_df.energy_needs_aMW)" + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "### Resource Generation Profiles\n", + "\n", + "We are looking for the least cost resource that can meet the energy need. We will make up some dummy generation profiles for wind and solar. " + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "def generate_resource_profiles(start_year=2024, years=10, noise_level=0.05):\n", + " # Generate a monthly date range\n", + " dates = pd.date_range(start=f\"{start_year}-01-01\", periods=years * 12, freq=\"M\")\n", + "\n", + " # Base profiles for wind and solar resources\n", + " wind_winter_peak = np.array(\n", + " [0.6, 0.55, 0.5, 0.45, 0.4, 0.35, 0.3, 0.25, 0.3, 0.4, 0.5, 0.55]\n", + " ) # Wind peaking in winter\n", + " wind_summer_peak = np.array(\n", + " [0.3, 0.25, 0.3, 0.4, 0.5, 0.55, 0.6, 0.55, 0.5, 0.45, 0.4, 0.35]\n", + " ) # Wind peaking in summer\n", + " solar_base = np.array(\n", + " [0.1, 0.2, 0.4, 0.6, 0.7, 0.9, 0.95, 0.85, 0.6, 0.4, 0.2, 0.1]\n", + " ) # Solar peaking in summer\n", + "\n", + " # Create dummy profiles with slight variations\n", + " profiles = {}\n", + " for i, location in enumerate([\"A\", \"B\", \"C\"]):\n", + " # Set a unique random seed for each location's wind and solar profiles\n", + " wind_seed = hash(f\"wind_{location}\") % (2**32)\n", + " solar_seed = hash(f\"solar_{location}\") % (2**32)\n", + "\n", + " np.random.seed(wind_seed)\n", + " if i == 0:\n", + " wind_profile = wind_winter_peak + noise_level * np.random.randn(12)\n", + " elif i == 1:\n", + " wind_profile = wind_summer_peak + noise_level * np.random.randn(12)\n", + " else:\n", + " # Create a mix or a balanced profile as the third option\n", + " wind_profile = 0.5 * (\n", + " wind_winter_peak + wind_summer_peak\n", + " ) + noise_level * np.random.randn(12)\n", + "\n", + " np.random.seed(solar_seed)\n", + " solar_profile = solar_base + noise_level * np.random.randn(12)\n", + "\n", + " # Set a unique random seed for the final noise adder\n", + " np.random.seed(wind_seed + 1)\n", + " profiles[f\"wind_{location}\"] = np.tile(\n", + " wind_profile, years\n", + " ) + noise_level * np.random.randn(len(dates))\n", + "\n", + " np.random.seed(solar_seed + 1)\n", + " profiles[f\"solar_{location}\"] = np.tile(\n", + " solar_profile, years\n", + " ) + noise_level * np.random.randn(len(dates))\n", + "\n", + " # Create a DataFrame with the profiles\n", + " df_profiles = pd.DataFrame(data=profiles, index=dates)\n", + "\n", + " return df_profiles" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "# Generate the dummy dataset\n", + "df_profiles = generate_resource_profiles()\n", + "\n", + "# Plot the profiles for the first 12 months\n", + "df_profiles.groupby(df_profiles.index.month).mean().plot()\n", + "plt.xlabel(\"month\")\n", + "plt.ylabel(\"normalized generation\")\n", + "plt.grid()" + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "As before, convert to xarray for linopy. " + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "df_profiles.index.name = \"datetime\"\n", + "df_profiles.columns.name = \"resource\"\n", + "resource_profiles_xr = xr.DataArray(df_profiles)" + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "### Resource Cost Profiles" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "def generate_unit_cost_profiles(start_year=2024, years=10, decay_rate=0.005):\n", + " # Generate a monthly date range\n", + " dates = pd.date_range(start=f\"{start_year}-01-01\", periods=years * 12, freq=\"M\")\n", + "\n", + " # Initial costs for wind and solar resources\n", + " initial_costs = {\n", + " \"wind_A\": 1.0,\n", + " \"wind_B\": 1.2,\n", + " \"wind_C\": 1.1,\n", + " \"solar_A\": 0.8,\n", + " \"solar_B\": 0.85,\n", + " \"solar_C\": 0.9,\n", + " }\n", + "\n", + " # Create dummy profiles with slow exponential decay\n", + " profiles = {}\n", + " for resource, initial_cost in initial_costs.items():\n", + " # Exponential decay function\n", + " decay = initial_cost * np.exp(-decay_rate * np.arange(len(dates)))\n", + "\n", + " # Assign decay curve to profile\n", + " profiles[resource] = decay\n", + "\n", + " # Create a DataFrame with the profiles\n", + " df_profiles = pd.DataFrame(data=profiles, index=dates)\n", + "\n", + " return df_profiles\n", + "\n", + "\n", + "# Generate the dummy dataset\n", + "df_unit_cost_profiles = generate_unit_cost_profiles()\n", + "\n", + "# Plot the unit cost profiles\n", + "df_unit_cost_profiles.plot()\n", + "plt.xlabel(\"dt\")\n", + "plt.ylabel(\"Unit Cost\")\n", + "plt.legend(title=\"resource\", bbox_to_anchor=(1.05, 1), loc=\"upper left\")\n", + "plt.grid()" + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "Convert to xarray" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "df_unit_cost_profiles.index.name = \"datetime\"\n", + "df_unit_cost_profiles.columns.name = \"resource\"\n", + "unit_cost_profiles_xr = xr.DataArray(df_unit_cost_profiles)" + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "## Optimized Energy Portfolio Model\n", + "\n", + "Now we are finally to the porfolio model. We know our resource needs, the generation of each resource, and the cost of each resource, we can write an optimization model to select the least cost portfolio that meets the energy need. \n", + "\n", + "The model is contained in the cells below. I've broken up the code across cells so you can get an idea of what the model components look like. Because the variables are vectorized across multiple dimensions, a single constraint is actually a whole array of constraints. " + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "# Initialize the model\n", + "m = lp.Model()" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "# build out\n", + "build_out = m.add_variables(\n", + " lower=0,\n", + " dims=resource_profiles_xr.dims,\n", + " coords=resource_profiles_xr.coords,\n", + " name=\"build_out\",\n", + ")\n", + "## make the build-out increase\n", + "m.add_constraints(build_out >= build_out.shift(datetime=1), name=\"build_out_increasing\")" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "# gen\n", + "## multiply for each resource\n", + "gen = resource_profiles_xr * build_out\n", + "## sum across resources\n", + "total_gen = gen.sum(dim=\"resource\")\n", + "\n", + "# Meet energy need\n", + "m.add_constraints(total_gen >= energy_needs_xr, name=\"meet_energy_need\")" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "# cost\n", + "build_month = build_out - build_out.shift(datetime=1)\n", + "cost = build_month * unit_cost_profiles_xr\n", + "total_cost = cost.sum()\n", + "total_cost" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "# objective\n", + "m.add_objective(total_cost)\n", + "m" + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "# Solve" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "m.solve()" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "sol = m.solution" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "# convert to a Pandas DataFrame for plotting\n", + "sol_df = (\n", + " sol.to_dataframe()\n", + " .reset_index()\n", + " .pivot_table(values=\"build_out\", columns=\"resource\", index=\"datetime\")\n", + ")\n", + "sol_df.plot()\n", + "plt.grid()\n", + "plt.title(\"Least Cost Energy Portfolio\")" + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "We see that a mix a resources are the least cost way to meet the average energy need. " + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "## Concusion\n", + "Linopy's vectorized variables and constraints allow for a concise representation of an optimization problen with multiple dimensions: in this case, a timeseries with a number of potential resources. " + ] + } + ], + "metadata": { + "kernelspec": { + "display_name": "qtenv", + "language": "python", + "name": "python3" + }, + "language_info": { + "codemirror_mode": { + "name": "ipython", + "version": 3 + }, + "file_extension": ".py", + "mimetype": "text/x-python", + "name": "python", + "nbconvert_exporter": "python", + "pygments_lexer": "ipython3", + "version": "3.11.3" + } + }, + "nbformat": 4, + "nbformat_minor": 2 +} diff --git a/examples/sudoku.ipynb b/examples/sudoku.ipynb new file mode 100644 index 00000000..b5230922 --- /dev/null +++ b/examples/sudoku.ipynb @@ -0,0 +1,493 @@ +{ + "cells": [ + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "# Sudoku Solver with Linopy\n", + "\n", + "The linopy package is good at tracking indices and writing optimization problems as vectorized functions (i.e., functions that operate on all rows of a column at once). This notebook adapts the Medium article [Creating Sudoku Solver with Python and Pyomo in Easy Steps](https://medium.com/@dhanalakotamohan314/creating-sudoku-solver-with-python-and-pyomo-in-easy-steps-fe22ec916090) by Dhanalakota Mohan for linopy to demonstrate how indices (i.e., dimensions and coordinates) work." + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "import matplotlib.pyplot as plt\n", + "import numpy as np\n", + "import pandas as pd\n", + "import xarray as xr\n", + "from mpl_toolkits.mplot3d.art3d import Poly3DCollection\n", + "\n", + "import linopy\n", + "\n", + "pd.options.display.float_format = \"{:.0f}\".format" + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "# The Puzzle\n", + "\n", + "Sudoku is a 9x9 grid where each of the digits 1 through 9 appears once and only once in each row, column, and 3x3 square. The puzzle starts with values already filled into some cells, which we call *hints*. Here is an example puzzle with its hints as a pandas DataFrame." + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "# look at the hints for the puzzle\n", + "puzzle_hints = [\n", + " (1, 7, 2),\n", + " (2, 2, 8),\n", + " (2, 6, 7),\n", + " (2, 8, 9),\n", + " (3, 1, 6),\n", + " (3, 3, 2),\n", + " (3, 7, 5),\n", + " (4, 2, 7),\n", + " (4, 5, 6),\n", + " (5, 4, 9),\n", + " (5, 6, 1),\n", + " (6, 5, 2),\n", + " (6, 8, 4),\n", + " (7, 3, 5),\n", + " (7, 7, 6),\n", + " (7, 9, 3),\n", + " (8, 2, 9),\n", + " (8, 4, 4),\n", + " (8, 8, 7),\n", + " (9, 3, 6),\n", + "]\n", + "puzzle_hints = pd.DataFrame(puzzle_hints, columns=[\"row\", \"column\", \"digit\"])\n", + "puzzle_hints_piv = puzzle_hints.pivot(\n", + " index=\"row\", columns=\"column\", values=\"digit\"\n", + ").replace(np.nan, \"-\")\n", + "puzzle_hints_piv" + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "# Defining Dimensions and Coordinates in Xarray\n", + "In Pandas parlance, the `index` are the row labels and `columns` are the column labels. When converting to Xarray, it is handy to name the index and columns with `df.index.name = \"row_name\"` and `df.columns.name = \"column_name\"`. \n", + "\n", + "In Xarray parlance, the Pandas `df.index.name` and `df.columns.name` are both called `dimensions`. Then the array of values of the index and columns are called `coordinates`. Think of latitude and longitude as `dimensions` with `coordinate` values. Xarray will infer the dimensions and coordinates from a Pandas DataFrame, provided columns and the index have names. The following cell shows how this looks. " + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "xr.DataArray(puzzle_hints_piv)" + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "In the Mohan article, the puzzle is solved with a clever reformulation. The obvious encoding — one integer variable per cell holding its digit — is hard to optimize: the rule \"the nine cells of a row are all different\" is not a linear constraint, and neither is \"this cell equals 5\".\n", + "\n", + "Instead we add a **third dimension**, `digit`, and make the variable a **binary on/off switch**: `sudoku[row, column, digit] == 1` means \"this cell holds this digit\". Every Sudoku rule then becomes a *linear* statement about sums of 0/1 values — for example, \"exactly one digit per cell\" is just \"the nine switches for that cell sum to 1\". That linearity is what a linear solver needs, and it is the whole reason for the extra dimension.\n", + "\n", + "Here is what an empty (all zeros) xarray DataArray of these switches looks like:" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "range_coord = range(1, 10)\n", + "xr.DataArray(\n", + " np.zeros(shape=(9, 9, 9)),\n", + " dims=[\"row\", \"column\", \"digit\"],\n", + " coords={\"row\": range_coord, \"column\": range_coord, \"digit\": range_coord},\n", + ")" + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "And here is a graphical representation of the 9x9x9 cube of switches — each small cube is one `(row, column, digit)` variable:" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "# Function to create a single cube at a specific location\n", + "def create_cube(x, y, z, size=1):\n", + " # Define the vertices of a cube\n", + " vertices = [\n", + " [x, y, z],\n", + " [x + size, y, z],\n", + " [x + size, y + size, z],\n", + " [x, y + size, z],\n", + " [x, y, z + size],\n", + " [x + size, y, z + size],\n", + " [x + size, y + size, z + size],\n", + " [x, y + size, z + size],\n", + " ]\n", + " # Define the 6 faces of the cube\n", + " faces = [\n", + " [vertices[j] for j in [0, 1, 2, 3]],\n", + " [vertices[j] for j in [4, 5, 6, 7]],\n", + " [vertices[j] for j in [0, 3, 7, 4]],\n", + " [vertices[j] for j in [1, 2, 6, 5]],\n", + " [vertices[j] for j in [0, 1, 5, 4]],\n", + " [vertices[j] for j in [2, 3, 7, 6]],\n", + " ]\n", + " return faces\n", + "\n", + "\n", + "# Create a 3D plot\n", + "fig = plt.figure(figsize=(7, 4))\n", + "ax = fig.add_subplot(111, projection=\"3d\")\n", + "\n", + "# Colors for the slices\n", + "colors = plt.cm.viridis(np.linspace(0, 1, 9))\n", + "\n", + "# Plot each cube\n", + "for i in range(9):\n", + " for j in range(9):\n", + " for k in range(9):\n", + " faces = create_cube(i, j, k)\n", + " poly3d = Poly3DCollection(\n", + " faces, color=colors[j], edgecolor=\"k\"\n", + " ) # Color by digit dimension\n", + " ax.add_collection3d(poly3d)\n", + "\n", + "# Set custom labels\n", + "ax.set_xlabel(\"Column\")\n", + "ax.set_ylabel(\"Digit\")\n", + "ax.set_zlabel(\"\")\n", + "ax.set_title(\"Sudoku Dimensions and Coordinates\")\n", + "\n", + "# Customize ticks and labels\n", + "ticks = np.arange(9)\n", + "labels = np.arange(1, 10) # Labels from 1 to 9\n", + "ax.set_xticks(ticks)\n", + "ax.set_yticks(ticks)\n", + "ax.set_zticks([]) # Remove z-axis ticks\n", + "ax.set_xticklabels(labels)\n", + "ax.set_yticklabels(labels)\n", + "\n", + "# Set the aspect ratio to be equal and adjust limits to fit the plot\n", + "ax.set_box_aspect([1, 1, 1]) # aspect ratio is 1:1:1\n", + "ax.set_xlim([0, 9])\n", + "ax.set_ylim([0, 9])\n", + "ax.set_zlim([0, 9])\n", + "\n", + "# Adjust layout to ensure labels are visible\n", + "plt.subplots_adjust(left=0.5, right=0.8, top=0.9, bottom=0.0)\n", + "\n", + "# Manually add z-axis labels\n", + "for z in range(9):\n", + " ax.text(x=-1, y=-1, z=z + 0.5, s=9 - z)\n", + "ax.text(x=-4, y=-1, z=4.5, s=\"Row\", ha=\"center\")\n", + "\n", + "plt.tight_layout()" + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "## Indices for the 3x3 Squares" + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "If you play Sudoku, you know each 3x3 square must also contain each digit once and only once. To express that, we first need to know which square every cell belongs to. We compute a square number for each `(row, column)` and store it as an xarray; later the model groups cells by this number to build the square constraints. Here is how the squares are numbered:" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "# Define the function to get the 3x3 square index\n", + "def get_square_index(row, col):\n", + " return (row // 3) * 3 + (col // 3)\n", + "\n", + "\n", + "# Create the grid with row, column, and square indices\n", + "square_index = pd.DataFrame(\n", + " {\n", + " \"row\": np.repeat(range_coord, 9),\n", + " \"column\": list(range_coord) * 9,\n", + " \"square\": [\n", + " get_square_index(row - 1, col - 1) + 1\n", + " for row in range_coord\n", + " for col in range_coord\n", + " ],\n", + " }\n", + ")\n", + "square_index = square_index.pivot(index=\"row\", columns=\"column\", values=\"square\")\n", + "square_index" + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "We put that into an xarray. Each cell now carries the number of the 3x3 square it belongs to:" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "square_index = xr.DataArray(square_index).rename(\"square\")\n", + "square_index" + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "# Model\n", + "\n", + "Now we build the model. Initialize a linopy model:" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "model = linopy.Model()" + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "Make the binary `sudoku` variable: a 9x9x9 cube of on/off switches over the `row`, `column`, and `digit` dimensions." + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "sudoku = model.add_variables(\n", + " name=\"sudoku\",\n", + " coords=[range_coord, range_coord, range_coord],\n", + " dims=[\"row\", \"column\", \"digit\"],\n", + " binary=True,\n", + ")" + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "It is handy to look at what linopy is doing. The code defines a single variable, but it is indexed on a 9x9x9 cube, so it really holds 729 binary variables." + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "sudoku" + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "Now the constraints. Each one sums the on/off switches over a single dimension and sets that sum to 1, which forces exactly one switch to be on along that dimension. The three lines below read as:\n", + "\n", + "- summing over **column** → each digit appears once in every **row**,\n", + "- summing over **row** → each digit appears once in every **column**,\n", + "- summing over **digit** → each **cell** holds exactly one digit.\n", + "\n", + "The last constraint is displayed by Jupyter, a handy way to see the structure linopy has built." + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "model.add_constraints(sudoku.sum(dim=[\"column\"]) == 1, name=\"row_digit_constraint\")\n", + "model.add_constraints(sudoku.sum(dim=[\"row\"]) == 1, name=\"column_digit_constraint\")\n", + "model.add_constraints(sudoku.sum(dim=[\"digit\"]) == 1, name=\"row_column_constraint\")" + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "Finally, the 3x3 square rule: each digit appears once and only once in every square. We `groupby` the cells by their `square_index` and sum within each group. Because the grouper is 2D (indexed by `row` and `column`) while the variable also carries `digit`, the grouped sum collapses `row` and `column` into a single `square` dimension and keeps `digit` — giving one constraint per `(square, digit)`, all in one vectorized call.\n", + "\n", + "> **Note:** `use_fallback=True` is a temporary workaround for a linopy bug ([#823](https://github.com/PyPSA/linopy/issues/823)) in the default grouped-sum path for multi-dimensional groupers. The fix is in [#824](https://github.com/PyPSA/linopy/pull/824); once it is merged, drop `use_fallback=True` and `.sum()` alone will be correct." + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "model.add_constraints(\n", + " sudoku.groupby(square_index).sum(use_fallback=True) == 1,\n", + " name=\"square_digit_constraint\",\n", + ")" + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "A Sudoku puzzle starts with some cells already filled in — the hints. We fix these too. Each hint is a `(row, column, digit)` triple; comparing each hint coordinate against the matching grid axis broadcasts a `True` wherever they line up, and `&` keeps only the cells that match on all three. Reducing with `.any(\"hint\")` unions the 20 hints into one boolean mask, which we hand to `add_constraints` to pin those switches to 1 — again in a single vectorized call." + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "hints = puzzle_hints.to_xarray().rename(index=\"hint\")\n", + "hint_mask = (\n", + " (sudoku.coords[\"row\"] == hints.row)\n", + " & (sudoku.coords[\"column\"] == hints.column)\n", + " & (sudoku.coords[\"digit\"] == hints.digit)\n", + ").any(\"hint\")\n", + "model.add_constraints(sudoku == 1, mask=hint_mask, name=\"hints\")" + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "This is really a *feasibility* problem: the constraints alone pin down the unique solution, so solving just means finding any point that satisfies them. An objective is optional here — we add a dummy one (maximize the sum of the switches) purely because it can nudge the solver; its value carries no meaning." + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "model.add_objective(sudoku.sum(), sense=\"max\")" + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "# Solve" + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "Solve the model with the default solver (HiGHS)." + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "model.solve()" + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "The solution comes out as an xarray representing a cube of binary variables. " + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "solution = model.solution.sudoku\n", + "solution" + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "Finally, turn the cube of 0/1 switches back into a 9x9 grid of digits. Since exactly one switch is on per cell, multiplying each switch by its own `digit` coordinate and summing over `digit` recovers the chosen digit — the exact inverse of the on/off encoding we started with." + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "final_sudoku_grid = (solution * solution.digit).sum(\"digit\").astype(int)\n", + "\n", + "result = pd.DataFrame(\n", + " data=final_sudoku_grid.values,\n", + " columns=puzzle_hints_piv.columns,\n", + " index=puzzle_hints_piv.index,\n", + ")\n", + "result" + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "Looks like it was solved! " + ] + } + ], + "metadata": { + "kernelspec": { + "display_name": "qtenv2", + "language": "python", + "name": "python3" + }, + "language_info": { + "codemirror_mode": { + "name": "ipython", + "version": 3 + }, + "file_extension": ".py", + "mimetype": "text/x-python", + "name": "python", + "nbconvert_exporter": "python", + "pygments_lexer": "ipython3", + "version": "3.11.3" + } + }, + "nbformat": 4, + "nbformat_minor": 2 +}