{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Multi-Level Disorder\n",
    "\n",
    "This guide demonstrates how to enumerate structures with disorder on **multiple independent subsets** of sites, such as simultaneous cation and anion disorder.\n",
    "\n",
    "## The Problem\n",
    "\n",
    "Sometimes you need to explore disorder on more than one subset of sites. For example, we might be interested in mixed cation/anion disorder in (Ti,Zr)O<sub>2</sub>:\n",
    "- **Cation disorder**: Ti/Zr substitution on metal sites\n",
    "- **Anion disorder**: O/F substitution on anion sites\n",
    "\n",
    "These disorders are independent—they occur on different sets of sites—but they interact through symmetry breaking.\n",
    "\n",
    "## The Hierarchical Approach\n",
    "\n",
    "When you disorder one subset of sites, the resulting structure typically has lower symmetry than the parent. This reduced symmetry should be used when enumerating disorder on the second subset.\n",
    "\n",
    "### Algorithm\n",
    "\n",
    "1. **Level 1**: Enumerate disorder on the first subset using the parent structure's symmetry\n",
    "2. **Level 2**: For each Level 1 configuration:\n",
    "   - The configuration has its own (typically reduced) symmetry\n",
    "   - Enumerate disorder on the second subset using this reduced symmetry\n",
    "3. Collect all Level 2 structures\n",
    "\n",
    "This hierarchical approach ensures:\n",
    "- No duplicate structures\n",
    "- Correct symmetry analysis at each level\n",
    "- Computational efficiency through symmetry reduction at each level"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Example: Ti/Zr and O/F Disorder in TiOF<sub>2</sub>\n",
    "\n",
    "Let us work through a realistic example: a 2×2×2 supercell of TiOF<sub>2</sub> with disorder on both cation and anion sublattices.\n",
    "\n",
    "### Setting Up the Parent Structure"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Created supercell with 32 atoms\n",
      "  - 8 Ti sites\n",
      "  - 0 X sites (to be O/F)\n"
     ]
    }
   ],
   "source": [
    "import numpy as np\n",
    "from pymatgen.core import Structure, Lattice\n",
    "from bsym.interface.pymatgen import unique_structure_substitutions\n",
    "\n",
    "# Create 2×2×2 TiOF2 supercell\n",
    "a = 3.798  # lattice parameter in Ångströms\n",
    "\n",
    "coords = np.array([[0.0, 0.0, 0.0],\n",
    "                   [0.5, 0.0, 0.0],\n",
    "                   [0.0, 0.5, 0.0],\n",
    "                   [0.0, 0.0, 0.5]])\n",
    "atom_list = ['Ti', 'X', 'X', 'X']\n",
    "lattice = Lattice.from_parameters(a=a, b=a, c=a, alpha=90, beta=90, gamma=90)\n",
    "unit_cell = Structure(lattice, atom_list, coords)\n",
    "\n",
    "# Create a 2×2×2 supercell\n",
    "parent_structure = unit_cell * [2, 2, 2]\n",
    "print(f\"Created supercell with {len(parent_structure)} atoms\")\n",
    "print(f\"  - {len([s for s in parent_structure if s.species_string == 'Ti'])} Ti sites\")\n",
    "print(f\"  - {len([s for s in parent_structure if s.species_string == 'X'])} X sites (to be O/F)\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Hierarchical Enumeration\n",
    "\n",
    "We will enumerate Ti/Zr disorder first (smaller combinatorial space), then O/F disorder for each Ti/Zr arrangement."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Level 1: Enumerating Ti/Zr arrangements...\n",
      "Found 3 unique Ti/Zr arrangements\n",
      "\n"
     ]
    }
   ],
   "source": [
    "# Level 1: Ti/Zr disorder on cation sites\n",
    "print(\"Level 1: Enumerating Ti/Zr arrangements...\")\n",
    "level1_structures = unique_structure_substitutions(\n",
    "    parent_structure,\n",
    "    'Ti',  # Substitute Ti sites\n",
    "    {'Ti': 6, 'Zr': 2}  # 6 Ti, 2 Zr\n",
    ")\n",
    "print(f\"Found {len(level1_structures)} unique Ti/Zr arrangements\\n\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Level 2: Enumerating O/F arrangements for each Ti/Zr configuration...\n",
      "  Ti/Zr config 1: 29371 O/F arrangements\n",
      "  Ti/Zr config 2: 28955 O/F arrangements\n",
      "  Ti/Zr config 3: 9782 O/F arrangements\n",
      "\n",
      "Total unique structures: 68108\n"
     ]
    }
   ],
   "source": [
    "# Level 2: O/F disorder for each Ti/Zr arrangement\n",
    "print(\"Level 2: Enumerating O/F arrangements for each Ti/Zr configuration...\")\n",
    "all_structures = []\n",
    "\n",
    "for i, structure in enumerate(level1_structures):\n",
    "    level2_structures = unique_structure_substitutions(\n",
    "        structure,  # Uses the reduced symmetry of this Ti/Zr arrangement\n",
    "        'X',\n",
    "        {'O': 8, 'F': 16}\n",
    "    )\n",
    "    print(f\"  Ti/Zr config {i+1}: {len(level2_structures)} O/F arrangements\")\n",
    "    all_structures.extend(level2_structures)\n",
    "\n",
    "print(f\"\\nTotal unique structures: {len(all_structures)}\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Understanding the Output\n",
    "\n",
    "In this example:\n",
    "- **Level 1** produces a modest number of unique Ti/Zr cation arrangements\n",
    "- **Level 2** varies for each Ti/Zr configuration:\n",
    "  - Some Ti/Zr arrangements preserve more symmetry → fewer O/F arrangements needed\n",
    "  - Other Ti/Zr arrangements break more symmetry → more O/F arrangements needed\n",
    "- The total combines all possibilities across both disorder types"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Implementation Pattern\n",
    "\n",
    "The general pattern for multi-level disorder is:\n",
    "\n",
    "```python\n",
    "# Level 1: First disorder type\n",
    "level1_configs = unique_structure_substitutions(\n",
    "    parent_structure, \n",
    "    species_to_substitute_1, \n",
    "    composition_1\n",
    ")\n",
    "\n",
    "# Level 2: Second disorder type\n",
    "all_configs = []\n",
    "for config in level1_configs:\n",
    "    level2_configs = unique_structure_substitutions(\n",
    "        config,  # Each has its own symmetry\n",
    "        species_to_substitute_2, \n",
    "        composition_2\n",
    "    )\n",
    "    all_configs.extend(level2_configs)\n",
    "```\n",
    "\n",
    "This pattern extends naturally to three or more levels by adding additional nested loops."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Iterative Approach for Multiple Levels\n",
    "\n",
    "For cases with three or more disorder types, nested loops become unwieldy. An iterative approach can be used instead that generalises to any number of levels:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {},
   "outputs": [],
   "source": [
    "def enumerate_multilevel_disorder(parent_structure, disorder_specs):\n",
    "    \"\"\"\n",
    "    Enumerate structures with multiple levels of disorder.\n",
    "    \n",
    "    Args:\n",
    "        parent_structure: Initial pymatgen Structure\n",
    "        disorder_specs: List of dicts, each containing:\n",
    "            - 'to_substitute': species label to replace\n",
    "            - 'site_distribution': dict of {species: count}\n",
    "    \n",
    "    Returns:\n",
    "        List of Structure objects with all disorder levels applied\n",
    "    \"\"\"\n",
    "    structures = [parent_structure]\n",
    "    \n",
    "    for level, spec in enumerate(disorder_specs, 1):\n",
    "        print(f\"Level {level}: Enumerating {spec['to_substitute']} → {spec['site_distribution']}\")\n",
    "        new_structures = []\n",
    "        \n",
    "        for i, structure in enumerate(structures):\n",
    "            if i % 100 == 0 and len(structures) > 100:\n",
    "                print(f\"  Processing structure {i+1}/{len(structures)}...\")\n",
    "            \n",
    "            level_structures = unique_structure_substitutions(\n",
    "                structure,\n",
    "                spec['to_substitute'],\n",
    "                spec['site_distribution']\n",
    "            )\n",
    "            new_structures.extend(level_structures)\n",
    "        \n",
    "        print(f\"  Generated {len(new_structures)} structures\\n\")\n",
    "        structures = new_structures\n",
    "    \n",
    "    return structures"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Level 1: Enumerating Ti → {'Ti': 6, 'Zr': 2}\n",
      "  Generated 3 structures\n",
      "\n",
      "Level 2: Enumerating X → {'O': 8, 'F': 16}\n",
      "  Generated 68108 structures\n",
      "\n",
      "Total unique structures: 68108\n"
     ]
    }
   ],
   "source": [
    "disorder_specs = [\n",
    "    {\n",
    "        'to_substitute': 'Ti',\n",
    "        'site_distribution': {'Ti': 6, 'Zr': 2}\n",
    "    },\n",
    "    {\n",
    "        'to_substitute': 'X',\n",
    "        'site_distribution': {'O': 8, 'F': 16}\n",
    "    }\n",
    "]\n",
    "\n",
    "# Run the multi-level enumeration\n",
    "all_structures = enumerate_multilevel_disorder(parent_structure, disorder_specs)\n",
    "\n",
    "print(f\"Total unique structures: {len(all_structures)}\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "This iterative approach:\n",
    "- Generalizes easily to 3+ disorder levels\n",
    "- Avoids deeply nested loops\n",
    "- Makes it easy to modify or reorder disorder specifications\n",
    "- Provides progress tracking for long enumerations"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Notes\n",
    "\n",
    "### Current Implementation\n",
    "\n",
    "This hierarchical enumeration can be implemented in two ways:\n",
    "\n",
    "1. **Manual chaining** (shown in the first example): Explicitly nest the `unique_structure_substitutions` calls. Best for 2 levels or when you need fine control over the process.\n",
    "\n",
    "2. **Iterative approach** (shown above): Use the `enumerate_multilevel_disorder` function to handle an arbitrary number of levels. Best for 3+ levels or when you want cleaner, more maintainable code.\n",
    "\n",
    "### Choosing the Level Order\n",
    "\n",
    "You can enumerate the subsets in any order—the final set of structures will be the same. However:\n",
    "- Starting with the **smaller combinatorial space** (fewer permutations) means fewer level-1 structures to loop over\n",
    "- Starting with disorder that **breaks symmetry more** may lead to faster level-2 enumerations\n",
    "- In practice, performance is often similar regardless of ordering\n",
    "\n",
    "When using the iterative approach, simply reorder the entries in `disorder_specs` to change the enumeration order.\n",
    "\n",
    "### Memory Considerations\n",
    "\n",
    "For very large structure sets:\n",
    "- Process structures in batches rather than storing all in memory\n",
    "- In the iterative approach, you can modify the function to write structures to disk after each level rather than keeping them all in the `structures` list"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Summary\n",
    "\n",
    "- **Multi-level disorder** requires enumerating permutations on multiple independent subsets of sites\n",
    "- **The hierarchical approach** enumerates one level at a time, using the appropriate symmetry at each step\n",
    "- **Two implementation patterns** are available:\n",
    "  - Manual chaining for 2 levels or fine-grained control\n",
    "  - Iterative approach for 3+ levels or cleaner code\n",
    "- Each level uses the symmetry of configurations from the previous level, ensuring correct and efficient enumeration"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3 (ipykernel)",
   "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.12.9"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}