Simulated annealing for magic square-of-squares puzzle

Question

I'm working on a function that takes in a NumPy array containing only mutually distinct positive square numbers. I want the function to pick an element at random, find if there is a lower positive square number than the number selected: (1) if there is, return the array with the lower positive square number in place of the original element (2) if there isn't, return a copy of the original array.

The reason I'm developing this function is to create a simulated annealing solution to the Magic Square of Squares Puzzle, and this function would serve as one of the mutations to the candidate solution.

I've had a go at the implementation:

def last_square(x):
    xcopy = x.copy()
    index_1 = np.random.randint(0, x.shape[0]), np.random.randint(0, x.shape[1])
    c = x[index_1]-1
    while np.sqrt(c) != int(np.sqrt(c)) or c in set(x.flatten()):
        if c < 2:
            return xcopy
        else:
            c -= 1
    x[index_1] = c
    return x

Will my function always have the behaviour I've specified?

Edit 1

I've been asked some clarifying questions that I'll address here. There were very helpful in getting me to see what else I still had to explain.

Why are you iterating through all possible numbers less than the chosen one, instead of just iterating through the array itself?

I am looking for the next smallest positive square number below the one that was randomly selected from the array that is not in the array.

Why do you care about sqrt in last_square at all? Taking the square root of the elements will not change their ordering.

I agree that removing the square too that not change the ordering of the the square numbers currently within the array, nor does it change the position of a candidate square number c among all positive square numbers.

I care about checking if np.sqrt(c) != int(np.sqrt(c)) because a number whose square root is an integer is a square number.

Do you care whether the lower number is chosen non-deterministically? What should happen if there is more than one lower number? Currently it looks like you've implemented a lookup for the next lowest number closest to the chosen number, which is not what you described in the question.

The randomized selection of an element in the array is of course random, but taking that element being chosen, I'm aiming for the rest of the behaviour of the function to be deterministic. That is to say, given the array and the randomly selected element, then the function will always give the same return.

I want to preserve the mutual distinctness of the positive square numbers within the array, so my goal is to find the next lowest positive square number that is not in the array. I was trying to acheive this by inclusively disjoining c in set(x.flatten()) to the header of the loop, thereby keeping the loop going when c is still a number within the array.

As written, if you find a lower number and copy it to the chosen location, your claim that the array's elements are mutually distinct will be violated.

Using the following

if c < 2:
            return xcopy
        else:
            c -= 1

I am checking if c has gone too low, in which case I return a copy of the original array. If I break from the loop, the following lines will change the original array and return the array with the substitution completed.

x[index_1] = c
    return x

I think what I am having trouble with in part may be when calling np.sqrt on c without checking if c is non-negative. I'll work on correcting that and editing the code when I've resolved that issue.

Let me know if there are further questions. I'm happy to clarify.

Answer 1

This is much more clear thanks to your edit.

Why do you store this as a list of squares? Can you store it as a unique list of positive integers, and square them later? That would make this massively simpler.

Another thing that would make this massively simpler is if you're able to guarantee that the list is maintained in some kind of order. In the example code below I do not assume this, so have to sort it myself. Maintaining order should be easy because the last_square mutation will not change order.

Will my function always have the behaviour I've specified?

Sounds like you need to write some unit tests :) You will have to mock your random function to make them predictable.

I care about checking if np.sqrt(c) != int(np.sqrt(c)) because a number whose square root is an integer is a square number.

Fine, but... you said that there was already a guarantee that these numbers are positive squares.

I tried to put together a vectorized version of this that does not loop. Rule number one when writing efficient Numpy code is to attempt to avoid loops. This solution roots, flattens and sorts the data, then looks for the closest gap. I do not know if it will perform better than yours, nor whether this is the most efficient implementation of this approach.

The performance will be influenced by the length of your data and the sparsity of your numbers. For highly sparse numbers the looping approach will probably do better. For highly dense numbers the approach below will probably do better.

def last_square_new(x):
    assert (x > 0).all()
    x = x.copy()

    # The n-dimensional random index. Cannot yet be used as an index.
    nd_index = np.random.default_rng().integers(
        np.zeros(x.ndim), x.shape
    )
    # What will the index of the chosen element be once the array is flattened?
    flat_index = nd_index[0]*x.shape[-1] + nd_index[1]
    # Make this usable as an index.
    nd_index = nd_index[:, np.newaxis].tolist()

    xsq = (np.sqrt(x) + 0.5).astype(np.uint64)     # Sqrt with safe conversion to int
    chosen = xsq[nd_index][0]                      # The randomly chosen element's sqrt
    xsq = xsq.flatten()                            # Flatten the root array
    s_indices = xsq.argsort()                      # Indices that would sort the root array
    flat_index = s_indices[s_indices][flat_index]  # Move the flat index to its sorted position
    xsq = xsq[s_indices][:flat_index+1]            # Sort the root array and truncate

    d_indices = np.arange(xsq.size-1)[np.diff(xsq) > 1]  # Which indices have gaps?
    if d_indices.size != 0:                              # Are there any gaps?
        gap_index = d_indices[-1]                        # Index of the closest gap (low side)
        best = xsq[gap_index+1] - 1                      # Best missing root
        x[nd_index] = best**2                            # Assign it to the output

    return x