/**
* Jupyter Notebook Generator for AutoRA Workflows
*
* Produces the same Python program as the Python code generator, but split
* into ordered notebook cells. The first code cell installs every package
* required by the imports that follow.
*
* @module utils/notebookGenerator
*/
import {
CodeBuilder,
TEMPLATES,
collectPipPackages,
flattenBlockNodes,
generateImports,
generateVariablesSetup,
generateWrapper,
prepareWorkflow
} from './pythonGenerator'
/**
* Build a notebook markdown cell from one or more lines of text.
*
* @param {string} text - Markdown source text.
* @returns {Object} Jupyter markdown cell object.
*/
function markdownCell(text) {
return {
cell_type: 'markdown',
metadata: {},
source: toSource(text)
}
}
/**
* Build a notebook code cell from a block of source text.
*
* @param {string} text - Python source text.
* @returns {Object} Jupyter code cell object.
*/
function codeCell(text) {
return {
cell_type: 'code',
execution_count: null,
metadata: {},
outputs: [],
source: toSource(text)
}
}
/**
* Convert a string into the array-of-lines form Jupyter uses for `source`,
* where every line except the last keeps its trailing newline.
*
* @param {string} text - Multi-line source text.
* @returns {string[]} Array of lines, each but the last ending in a newline.
*/
function toSource(text) {
const lines = text.split('\n')
return lines.map((line, i) => (i < lines.length - 1 ? `${line}\n` : line))
}
/**
* Build the "run the workflow" code cell. Unlike the standalone Python file,
* the notebook runs the loop at top level so each cycle's state is inspectable.
*
* @param {Block[]} blocks - Ordered execution block tree (see getExecutionOrder in pythonGenerator; `Block` is defined there).
* @param {Map} componentMeta - Map of node id to component metadata (`varName`, `nodeName`, `pythonName`).
* @returns {string} Python source for the run cell.
*/
function buildRunCell(blocks, componentMeta) {
const code = new CodeBuilder()
// Variable setup + state initialization (templates are indented for a
// function body, so strip the leading 4 spaces for top-level use).
code.multiline(dedent(generateVariablesSetup(componentMeta, flattenBlockNodes(blocks))))
code.blank()
code.multiline(dedent(TEMPLATES.initState))
code.blank()
// Emit a single `state = fn(state[, num_samples=…])` call at the given level.
const addComponentCall = (node, level) => {
const meta = componentMeta.get(node.id)
if (!meta) return
const { varName, nodeName, pythonName } = meta
const isSampler = pythonName.includes('sample') || pythonName.includes('sampler')
code.indent(`# ${nodeName}`, level)
if (isSampler && node.parameters?.num_samples != null) {
code.indent(`state = ${varName}(state, num_samples=${node.parameters.num_samples})`, level)
} else {
code.indent(`state = ${varName}(state)`, level)
}
code.blank()
}
// Emit the block tree in order (see getExecutionOrder). At top level `once`
// blocks run unindented; `loop` blocks wrap their children in a for-loop and
// recurse, so nested loops render as nested for-loops. A loop prints its cycle
// only when it directly runs components (not when it only holds nested loops).
const renderBlocks = (blks, level, loopDepth = 0) => {
blks.forEach(block => {
if (block.type === 'loop') {
const loopVar = `cycle_${loopDepth}`
code.indent(`# Experiment loop (${block.maxCounter} cycles)`, level)
code.indent(`for ${loopVar} in range(${block.maxCounter}):`, level)
if (block.children.some(c => c.type === 'once')) {
code.indent(`print(f'Cycle {${loopVar}}')`, level + 1)
code.blank()
}
renderBlocks(block.children, level + 1, loopDepth + 1)
code.blank()
} else {
block.nodes.forEach(node => addComponentCall(node, level))
}
})
}
renderBlocks(blocks, 0)
code.line('print("Workflow completed!")')
code.line('state')
return code.toString()
}
/**
* Remove a single level (4 spaces) of leading indentation from each line so
* function-body templates can be reused at notebook top level.
*
* @param {string} text - Indented multi-line source text.
* @returns {string} Text with one 4-space indent level removed from each line.
*/
function dedent(text) {
return text
.split('\n')
.map(line => (line.startsWith(' ') ? line.slice(4) : line))
.join('\n')
}
/**
* Generate a Jupyter notebook (parsed JSON object) from workflow state.
*
* @param {Object} state - Editor state with `nodes`, `connections` and `components`.
* @returns {Object} Jupyter notebook object with `cells`, `metadata`, `nbformat` and `nbformat_minor`.
*/
export function generateNotebook(state) {
const { blocks, imports, componentMeta, derivesVariablesFromRunner, needsEquationVariables } = prepareWorkflow(state)
const cells = []
// 1. Title
cells.push(markdownCell(
`# AutoRA Workflow\n\n` +
`Generated by AutoRA Workflow Editor on ${new Date().toISOString()}`
))
// 2. Install dependencies (must cover every import that follows)
cells.push(markdownCell('## 1. Install dependencies'))
const pipPackages = collectPipPackages(state)
const pipCell = pipPackages.length === 0
? '# No additional packages required'
: `%pip install ${pipPackages.join(' ')}`
cells.push(codeCell(pipCell))
// 3. Imports
cells.push(markdownCell('## 2. Imports'))
const importsCode = new CodeBuilder()
generateImports(importsCode, imports, { usesPlaceholderVariables: !derivesVariablesFromRunner, usesEquationVariables: needsEquationVariables })
cells.push(codeCell(importsCode.toString()))
// 4. Component definitions (one code cell per component for clarity).
// Components that produce an identical definition (same function name and
// parameters) are emitted only once and reused, so no duplicate cells appear.
cells.push(markdownCell('## 3. Component definitions'))
const seenWrappers = new Set()
componentMeta.forEach((meta) => {
const wrapper = new CodeBuilder()
generateWrapper(wrapper, meta)
// Drop the trailing blank line each wrapper appends.
const src = wrapper.toString().replace(/\n+$/, '')
if (seenWrappers.has(src)) return
seenWrappers.add(src)
cells.push(codeCell(src))
})
// 5. Run the workflow
cells.push(markdownCell('## 4. Run the workflow'))
cells.push(codeCell(buildRunCell(blocks, componentMeta)))
return {
cells,
metadata: {
kernelspec: {
display_name: 'Python 3',
language: 'python',
name: 'python3'
},
language_info: {
name: 'python'
}
},
nbformat: 4,
nbformat_minor: 5
}
}
/**
* Generate the notebook as a formatted JSON string ready to write to a
* `.ipynb` file.
*
* @param {Object} state - Editor state with `nodes`, `connections` and `components`.
* @returns {string} Notebook serialized as a formatted JSON string.
*/
export function generateNotebookString(state) {
return JSON.stringify(generateNotebook(state), null, 1)
}